Metadata-Version: 2.5
Name: orionmoe
Version: 1.0.0
Summary: OrionMoE: A Mixture-of-Experts Foundation Model for Tabular In-Context Learning
Project-URL: Documentation, https://github.com/Lexsi-Labs/OrionMoE#README
Project-URL: Issues, https://github.com/Lexsi-Labs/OrionMoE/issues
Project-URL: Model weights, https://huggingface.co/Lexsi/OrionMoE
Project-URL: Homepage, https://www.lexsi.ai/
Author: Mohamed Bouadi, Nassim Bouarour, Aditya Tanna, Vinay Kumar Sankarapu
License-Expression: LicenseRef-LSAL-1.2-OrionMoE
License-File: LICENSE
License-File: NOTICE
Keywords: ICL,In-Context Learning,foundation model,mixture of experts,tabular data
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering
Requires-Python: <3.13,>=3.9
Requires-Dist: huggingface-hub<2,>=0.27
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: scikit-learn<1.9,>=1.3.0
Requires-Dist: scipy
Requires-Dist: torch<3,>=2.4
Description-Content-Type: text/markdown

# OrionMoE

<div align="center">
  <a href="https://www.lexsi.ai/">
    <img src="https://img.shields.io/badge/Lexsi-Homepage-FF6B6B?style=for-the-badge" alt="Homepage"/>
  </a>
  <a href="https://huggingface.co/Lexsi/OrionMoE">
    <img src="https://img.shields.io/badge/🤗%20Hugging%20Face-Lexsi%20AI-FFD21E?style=for-the-badge" alt="Hugging Face"/>
  </a>
  <br/>
  <a href="https://www.python.org/downloads/">
    <img src="https://img.shields.io/badge/python-3.9+-blue.svg" alt="Python 3.9+"/>
  </a>
  <a href="https://pytorch.org/">
    <img src="https://img.shields.io/badge/PyTorch-2.4+-orange.svg" alt="PyTorch"/>
  </a>
  <a href="https://huggingface.co/Lexsi/OrionMoE/blob/main/LICENSE">
      <img src="https://img.shields.io/badge/License-LSAL_1.2-blue.svg"
           alt="License: LSAL v1.2, OrionMoE Edition"/>
    </a>
</div>


OrionMoE is a mixture-of-experts foundation model for tabular classification. It learns from synthetic datasets and uses labeled examples as context to predict classes on a new dataset, without updating the model weights.

The model combines distribution-aware column embeddings, induced row attention, feature attention, and dataset-conditioned expert routing. Its scikit-learn interface supports numerical and categorical features, missing values, and classification beyond the checkpoint's native class count through many-class inference.

# Quick Start 

## Installation

Install the inference package from PyPI:

```bash
python -m pip install orionmoe
```

The package includes the sklearn wrapper, model architecture, preprocessing, and inference utilities. Model weights are downloaded separately from Hugging Face.

Python 3.9-3.12 and PyTorch 2.4 or newer are required. For GPU inference, install a CUDA-compatible PyTorch build using the [PyTorch installation guide](https://pytorch.org/get-started/locally/). CPU inference is also supported.

## Inference

The classifier resolves its checkpoint when `fit` is first called. With no `checkpoint_path`, it checks the Hugging Face cache and downloads the selected checkpoint if needed.

```python
from orion_moe.sklearn import OrionMoEClassifier

classifier = OrionMoEClassifier()
```

Set `allow_auto_download=False` for cache-only loading without network access. An existing `checkpoint_path` always takes precedence and does not access the hugging face Hub.

A complete example using a local checkpoint:

```python
import torch
from sklearn.datasets import load_breast_cancer
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.model_selection import train_test_split

from orion_moe.sklearn import OrionMoEClassifier

X, y = load_breast_cancer(return_X_y=True, as_frame=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42
)

classifier = OrionMoEClassifier(
    checkpoint_path="checkpoints/orionmoe-small-classifier.ckpt",
    device="cuda" if torch.cuda.is_available() else "cpu",
    ensemble_contexts=4,
    random_state=42,
)
classifier.fit(X_train, y_train)

probabilities = classifier.predict_proba(X_test)
predictions = classifier.classes_[probabilities.argmax(axis=1)]

print("Accuracy:", accuracy_score(y_test, predictions))
print("ROC AUC:", roc_auc_score(y_test, probabilities[:, 1]))
```

`fit` prepares preprocessing and stores labeled context; it does not fine-tune the checkpoint. `predict` returns class labels, and `predict_proba` returns probabilities with columns ordered as `classifier.classes_`.

Inputs can be NumPy arrays or pandas DataFrames. Keep the same feature columns and order at prediction time. The classifier handles preprocessing internally. For categorical columns stored as numeric codes, pass their names or indices through `categorical_features`.

## Inference settings

OrionMoE provides a set of parameters for customizing inference behavior. The main options below show the wrapper defaults:

```python
from orion_moe.sklearn import OrionMoEClassifier

classifier = OrionMoEClassifier(
    ensemble_contexts=8,          # Context/preprocessing members; more increases inference cost
    preprocessing_transforms=(
        "standard", "power_softclip", "quantile_uniform", "quantile_rtdl"
    ),                            # Transforms assigned across context members
    feat_shuffle_method="latin", # Feature permutation strategy
    class_shuffle_method="shift", # Class permutation strategy
    softmax_temperature=0.9,      # Lower values produce sharper class probabilities
    average_logits=False,        # Average probabilities; True averages logits instead
    categorical_features=None,   # Optional categorical column names or indices
    max_context_size=None,       # Use all labeled rows; set a limit to reduce context memory
    query_chunk_size=0,          # No explicit chunk size; e.g. 128 processes queries in chunks
    chunking_strategy="support_cached",  # Reuse support computations across query chunks
    oom_query_chunk_retry=True,  # Retry CUDA OOM failures with smaller query chunks
    oom_adaptive_policy=True,    # Adapt recovery settings to the current input
    oom_context_retry=True,      # Allow fewer labeled context rows during OOM recovery
    use_kv_cache=False,          # Enable attention KV caching
    persistent_kv_cache=False,   # Reuse KV caches across calls when caching is enabled
    checkpoint_path=None,       # Use the Hugging Face cache or download when needed
    allow_auto_download=True,    # Permit downloading a missing checkpoint
    checkpoint_version="orionmoe-small-classifier.ckpt",
    device=None,                # Automatically select CUDA or CPU
    dtype="float32",            # Also supports "bfloat16" and "float16" for CUDA AMP
    use_amp=False,               # CUDA mixed precision when dtype is float16/bfloat16
    random_state=0,             # Seed for context sampling and ensemble permutations
)
```

For smaller inference runs, reduce `query_chunk_size` and optionally set `max_context_size`. Reducing context changes which labeled examples are used. Automatic OOM recovery is enabled by default and can reduce query chunks, features, or context. Set `oom_feature_retry=False` and `oom_context_retry=False` when evaluation must preserve the chosen features and context.


## License

The OrionMoE source code and model weights are licensed under the **Lexsi Labs Source Available License (LSAL) v1.2, OrionMoE Edition**. This is a source-available license, not an OSI-approved open-source license.

- Personal noncommercial research and academic research and teaching are permitted subject to the license.
- Organizational use is subject to the written acknowledgement or permission requirements in Section 1A.
- Commercial use requires a separate commercial license.
- Production or user-facing deployment requires express written authorization.

See [LICENSE](LICENSE) for the complete terms. Copyright (c) 2025-2026 Lithasa Technologies Pvt. Ltd. (Lexsi Labs). For permissions and commercial licensing, contact **support@lexsi.ai**


## Contact

For questions, issues, or contributions, please:
- Open an issue on [GitHub](https://github.com/Lexsi-Labs/OrionMoE/issues)