Views
No views yet
1pip install uv
2uv pip install autogluon.tabular[mitra] 1import pandas as pd
2from autogluon.tabular import TabularDataset, TabularPredictor
3from sklearn.model_selection import train_test_split
4from sklearn.datasets import load_wine
5
6# Load datasets
7wine_data = load_wine()
8wine_df = pd.DataFrame(wine_data.data, columns=wine_data.feature_names)
9wine_df['target'] = wine_data.target
10
11print("Dataset shapes:")
12print(f"Wine: {wine_df.shape}")
13
14# Create train/test splits (80/20)
15wine_train, wine_test = train_test_split(wine_df, test_size=0.2, random_state=42, stratify=wine_df['target'])
16
17print("Training set sizes:")
18print(f"Wine: {len(wine_train)} samples")
19
20# Convert to TabularDataset
21wine_train_data = TabularDataset(wine_train)
22wine_test_data = TabularDataset(wine_test)
23
24# Create predictor with Mitra
25print("Training Mitra classifier on classification dataset...")
26mitra_predictor = TabularPredictor(label='target')
27mitra_predictor.fit(
28 wine_train_data,
29 hyperparameters={
30 'MITRA': {'fine_tune': False}
31 },
32 )
33
34print("\nMitra training completed!")
35
36# Make predictions
37mitra_predictions = mitra_predictor.predict(wine_test_data)
38print("Sample Mitra predictions:")
39print(mitra_predictions.head(10))
40
41# Show prediction probabilities for first few samples
42mitra_predictions = mitra_predictor.predict_proba(wine_test_data)
43print(mitra_predictions.head())
44
45# Show model leaderboard
46print("\nMitra Model Leaderboard:")
47mitra_predictor.leaderboard(wine_test_data)1mitra_predictor_ft = TabularPredictor(label='target')
2mitra_predictor_ft.fit(
3 wine_train_data,
4 hyperparameters={
5 'MITRA': {'fine_tune': True, 'fine_tune_steps': 10}
6 },
7 time_limit=120, # 2 minutes
8 )
9
10print("\nMitra fine-tuning completed!")
11
12# Show model leaderboard
13print("\nMitra Model Leaderboard:")
14mitra_predictor_ft.leaderboard(wine_test_data)@article{zhang2025mitra,
title={Mitra: Mixed synthetic priors for enhancing tabular foundation models},
author={Zhang, Xiyuan and Maddix, Danielle C and Yin, Junming and Erickson, Nick and Ansari, Abdul Fatir and Han, Boran and Zhang, Shuai and Akoglu, Leman and Faloutsos, Christos and Mahoney, Michael W and others},
journal={arXiv preprint arXiv:2510.21204},
year={2025}
}