Improve model card: Add pipeline tag, paper/project links, authors & full content

README.md

@@ -1,7 +1,12 @@
 ---
 license: mit
+pipeline_tag: table-question-answering
 ---
 
+This model is presented in the paper [Orion-Bix: Bi-Axial Attention for Tabular In-Context Learning](https://huggingface.co/papers/2512.00181).
+Authors: Mohamed Bouadi, Pratinav Seth, Aditya Tanna, Vinay Kumar Sankarapu
+Project Page: https://www.lexsi.ai/
+
 <div align="center">
   <img src="logo.png" alt="Orion-BiX Logo" width="700"/>
 </div>
@@ -10,7 +15,7 @@ license: mit
   <a href="https://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">
+  <a href="https://huggingface.co/Lexsi/Orion-BiX">
     <img src="https://img.shields.io/badge/🤗%20Hugging%20Face-Lexsi AI-FFD21E?style=for-the-badge" alt="Hugging Face"/>
   </a>
   <a href="https://discord.gg/dSB62Q7A">
@@ -21,52 +26,58 @@ license: mit
   </a>
 </div>
 
+[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-2.2+-orange.svg)](https://pytorch.org/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 
+# Orion-BiX: Bi-Axial Meta-Learning for Tabular In-Context Learning
 
-# Orion-BiX: Bi-Axial Meta-Learning Model for Tabular In-Context Learning
-
-**Orion-BiX** is an advanced tabular foundation model that combines **Bi-Axial Attention** with **Meta-Learning** capabilities for few-shot tabular classification. The model extends the TabICL architecture with alternating attention patterns and episode-based training.
+**[Orion-BiX](https://arxiv.org/abs/2512.00181)** is an advanced tabular foundation model that combines **Bi-Axial Attention** with **Meta-Learning** capabilities for few-shot tabular classification. The model extends the TabICL architecture with alternating attention patterns and episode-based training, achieving state-of-the-art performance on domain-specific benchamrks such as Healthcare and Finance.
 
-The model is part of **Orion**, a family of tabular foundation models with various enhancements.
+## 🏗️ Approach and Architecture
 
 ### Key Innovations
 
-1. **Bi-Axial Attention**: Alternating attention patterns (Standard → Grouped → Hierarchical → Relational) that capture multi-scale feature interactions within tabular data
-2. **Meta-Learning with k-NN Support Selection**: Episode-based training with intelligent support set selection using similarity metrics
-3. **Three-Component Architecture**: Column embedding (Set Transformer), Bi-Axial row interaction, and In-Context Learning prediction
+Orion-BiX introduces three key architectural innovations:
+
+1.  **Bi-Axial Attention**: Alternating attention patterns (Standard → Grouped → Hierarchical → Relational) that capture multi-scale feature interactions
+2.  **Meta-Learning**: Episode-based training with k-NN support selection for few-shot learning
+3.  **Configurable Architecture**: Flexible design supporting various attention mechanisms and training modes
+4.  **Production Ready**: Memory optimization, distributed training support, and scikit-learn interface
+
+### Component Details
 
-### Architecture Overview
+Orion-BiX follows a three-component architecture:
 
 ```
-Input → tf_col (Set Transformer) → Bi-Axial Attention → tf_icl (ICL) → Output
+Input → Column Embedder (Set Transformer) → Bi-Axial Attention → ICL Predictor → Output
 ```
-**Component Details:**
-- **tf_col (Column Embedder)**: Set Transformer for statistical distribution learning across features
-- **Bi-Axial Attention**: Replaces standard RowInteraction with alternating attention patterns:
-  - Standard Cross-Feature Attention
-  - Grouped Feature Attention
-  - Hierarchical Feature Attention
-  - Relational Feature Attention
-  - CLS Token Aggregation
-- **tf_icl (ICL Predictor)**: In-context learning module for few-shot prediction
 
+1.  **Column Embedder**: Set Transformer for statistical distribution learning across features from TabICL
+2.  **Bi-Axial Attention**: Replaces standard RowInteraction with alternating attention patterns:
+    -   **Standard Cross-Feature Attention**: Direct attention between features
+    -   **Grouped Feature Attention**: Attention within feature groups
+    -   **Hierarchical Feature Attention**: Hierarchical feature patterns
+    -   **Relational Feature Attention**: Full feature-to-feature attention
+    -   **CLS Token Aggregation**: Multiple CLS tokens (default: 4) for feature summarization
+3.  **tf_icl ICL Predictor**: In-context learning module for few-shot prediction
 
-## Usage
-
-```python
-from orion_bix.sklearn import OrionBiXClassifier
+Each `BiAxialAttentionBlock` applies four attention patterns in sequence:
 
-# Initialize and use
-clf = OrionBiXClassifier()
-clf.fit(X_train, y_train)
-predictions = clf.predict(X_test)
 ```
-
-This code will automatically download the pre-trained model from Hugging Face and use a GPU if available.
+Standard → Grouped → Hierarchical → Relational → CLS Aggregation
+```
 
 ## Installation
 
+### Prerequisites
+
+-   Python 3.9-3.12
+-   PyTorch 2.2+ (with CUDA support recommended)
+-   CUDA-capable GPU (recommended for training)
+
 ### From the source
+
 #### Option 1: From the local clone
 
 ```bash
@@ -75,19 +86,98 @@ pip install -e .
 ```
 
 #### Option 2: From the Git Remote
-
 ```bash
 pip install git+https://github.com/Lexsi-Labs/Orion-BiX.git
 ```
 
+## Usage
+
+Orion-BiX provides a scikit-learn compatible interface for easy integration:
+
+```python
+from orion_bix.sklearn import OrionBixClassifier
+
+# Initialize and fit the classifier
+clf = OrionBixClassifier()
+
+# Fit the model (prepares data transformations)
+clf.fit(X_train, y_train)
+
+# Make predictions
+predictions = clf.predict(X_test)
+probabilities = clf.predict_proba(X_test)
+```
+
+## Preprocessing
+
+Orion-BiX includes automatic preprocessing that handles:
+
+1.  **Categorical Encoding**: Automatically encodes categorical features using ordinal encoding
+2.  **Missing Value Imputation**: Handles missing values using median imputation for numerical features
+3.  **Feature Normalization**: Supports multiple normalization methods:
+    -   `"none"`: No normalization
+    -   `"power"`: Yeo-Johnson power transform
+    -   `"quantile"`: Quantile transformation to normal distribution
+    -   `"quantile_rtdl"`: RTDL-style quantile transform
+    -   `"robust"`: Robust scaling using median and quantiles
+4.  **Outlier Handling**: Clips outliers beyond a specified Z-score threshold (default: 4.0)
+5.  **Feature Permutation**: Applies systematic feature shuffling for ensemble diversity:
+    -   `"none"`: Original feature order
+    -   `"shift"`: Circular shifting
+    -   `"random"`: Random permutation
+    -   `"latin"`: Latin square patterns (recommended)
+
+The preprocessing is automatically applied during `fit()` and `predict()`, so no manual preprocessing is required.
+
+## Performance
+
+<div align="center">
+  <img src="figures/accuracy_ranking_talent.png" alt="Accuracy Ranking TALENT" width="700"/>
+</div>
+
+<div align="center">
+  <img src="figures/accuracy_ranking_tabzilla.png" alt="Accuracy Ranking TabZilla" width="700"/>
+</div>
+
+<div align="center">
+  <img src="figures/accuracy_ranking_openml-cc18.png" alt="Accuracy Ranking OPENML-CC18" width="700"/>
+</div>
+
+<div align="center">
+  <table>
+    <tr>
+      <td style="padding: 5px;"><img src="figures/relative_acc_improvement_over_tabzilla.png" alt="Relative Improvement over XGBoost on TabZilla" width="700"/></td>
+    </tr>
+  </table>
+</div>
+
 ## Citation
 
-If you use Orion-BiX in your research, please cite:
+If you use Orion-BiX in your research, please cite our [paper](https://arxiv.org/abs/2512.00181):
 
 ```bibtex
-@misc{bouadi25oriobix,
-  title={Orion-BiX: Bi-Axial Meta-Learning for Tabular In-Context Learning},
-  author={Mohamed Bouadi and Pratinav Seth and Aditya Tanna and Vinay Kumar Sankarapu},
-  year={2025},
+@article{bouadi2025orionbix,
+      title={Orion-Bix: Bi-Axial Attention for Tabular In-Context Learning}, 
+      author={Mohamed Bouadi and Pratinav Seth and Aditya Tanna and Vinay Kumar Sankarapu},
+      year={2025},
+      eprint={2512.00181},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2512.00181}, 
 }
-```
+```
+
+## License
+
+This project is released under the MIT License. See [LICENSE](LICENSE) for details.
+
+## Contact
+
+For questions, issues, or contributions, please:
+- Open an issue on [GitHub](https://github.com/Lexsi-Labs/Orion-BiX/issues)
+- Join our [Discord](https://discord.gg/dSB62Q7A) community
+
+
+## 🙏 Acknowledgments
+
+Orion-BiX is built on top of [TabICL](https://github.com/soda-inria/tabicl), a tabular foundation model for in-context learning. We gratefully acknowledge the TabICL authors for their foundational work and for making their codebase publicly available.