Skip to content

API Reference

RAGImputer

The main sklearn-compatible imputer. Combines correlation-weighted retrieval with LLM-based generation to impute missing values in tabular data.

tabrag_xai_imputer.RAGImputer

Bases: BaseEstimator, TransformerMixin

Retrieval-Augmented Generation Imputer.

Parameters:

Name Type Description Default
n_neighbors int

Number of nearest neighbours to retrieve.

5
feature_weighting ('correlation', 'uniform')

How to weight observed features during retrieval. - "correlation" – weight by |correlation| with missing features. - "uniform" – equal weight for all observed features.

"correlation"
llm_model_name str

LLM model identifier.

"openai/gpt-4.1-nano"
llm_api ('open_router', 'gemini', 'gpt', 'claude')

API backend for LLM calls.

"open_router","gemini","gpt","claude"
dataset_name str

Injected into the LLM prompt.

"Unknown Dataset"
llm_batch_size int

Rows per LLM call.

1
min_complete_rows int

Minimum complete rows required in training data.

1
api_key str or None

API key for the chosen LLM backend. When None, the key is read from the environment variable (API_KEY_GEMINI, API_KEY_GPT, API_KEY_OPEN_ROUTER, or API_KEY_CLAUDE) via a .env file.

None

fit(X, y=None)

transform(X, y=None)

fit_transform(X, y=None, **fit_params)

explain(X_missing, X_imputed, row_indices=None)

Return LLM-generated explanations for each imputed row.

Parameters:

Name Type Description Default
X_missing array-like of shape (n_samples, n_features)

Original data containing NaN where values were missing.

required
X_imputed array-like of shape (n_samples, n_features)

Fully imputed data returned by transform.

required
row_indices list[int] | None

Indices of rows to explain. When None, all rows that had at least one missing value in X_missing are explained.

None

Returns:

Type Description
list[str]

One explanation string per requested row, in the same order as row_indices (or the auto-detected missing rows).

get_feature_names_out(input_features=None)