TableMage   🧙‍♂️📊

TableMage is a Python package for low-code/conversational clinical data science. TableMage can help you quickly explore tabular datasets, easily perform regression analyses, and effortlessly benchmark machine learning models.

Installation

We recommend installing TableMage in a new virtual environment.

To install TableMage:

git clone https://github.com/ajy25/TableMage.git
cd TableMage
pip install .

TableMage supports Python versions 3.10 through 3.12.

[!NOTE] For MacOS users: You might run into an error involving XGBoost, one of TableMage's dependencies, when using TableMage for the first time. To resolve this error, you'll need to install libomp: brew install libomp. This requries Homebrew.

Quick start (low-code)

You'll likely use TableMage for machine learning model benchmarking. Here's how to do it.

import tablemage as tm
import pandas as pd
import joblib

# load table (assume 'y' is a numeric variable we wish to predict)
df = ...

# initialize an Analyzer object
analyzer = tm.Analyzer(df, test_size=0.2)

# preprocess data, taking care to exclude the target variable 'y' from the operations
c

# train regressors
reg_report = analyzer.regress(  # categorical variables are automatically one-hot encoded
    models=[                    # hyperparameter tuning is preset and automatic
        tm.ml.LinearR('l2', name='ridge'),
        tm.ml.TreesR('random_forest', name='rf'),
        tm.ml.TreesR('xgboost', name='xgb'),
    ],
    target='y',                 # automatically drops examples with missing values in target variable
    predictors=None,            # None signifies all variables except target variable
    feature_selectors=[
        tm.fs.BorutaFSR()       # select subset of predictors prior to training
    ]
)

# view model metrics
print(reg_report.metrics('test'))

# predict on new data
new_df = ...
ridge_model = reg_report.model('ridge').sklearn_pipeline()
y_pred = ridge_model.predict(new_df)

# save as sklearn pipeline
joblib.dump(ridge_model, 'ridge.joblib')

Quick start (conversational)

First, install the required additional dependencies.

pip install '.[agents]'

Next, add your API key. You only need to do this once; your API key will be written to a local .env file.

import tablemage as tm
tm.use_agents()                                             # import the agents module
tm.agents.set_key("openai", "add-your-api-key-here")        # set API key

You can open up a chat user interface by running the following code and navigating to the URL that appears in the terminal. Your conversation with the ChatDA, the AI agent, appears on the left, while ChatDA's analyses (figures made, tables produced, TableMage commands used) appear on the right.

import tablemage as tm
tm.use_agents()
tm.agents.options.set_llm(
    llm_type="openai", 
    model_name="gpt-4o-mini", 
    temperature=0.1
)
# optionally, multimodal ChatDA can interpret figures
tm.agents.options.set_multimodal_llm(
    llm_type="openai",
    model_name="gpt-4o-mini",
    temperature=0.1
)                           # multimodal LLM must be specified for multimodal ChatDA
tm.agents.App(
    multimodal=True         # additional parameters can be set, e.g. memory type, 
).run(debug=False)          # disabling/enabling Python environment, etc.

Or, you can chat with the AI agent directly in Python:

import pandas as pd
import tablemage as tm
tm.use_agents()
tm.agents.options.set_llm(
    llm_type="openai", 
    model_name="gpt-4o-mini", 
    temperature=0.1
)

# load table
df = ...

# initialize a ChatDA object
agent = tm.agents.ChatDA(
    df,                     # additional parameters can be set, e.g. memory type, 
    test_size=0.2           # disabling/enabling Python environment, etc.
)

# chat with the agent
print(agent.chat("Compute the summary statistics for the numeric variables."))

[!NOTE] You must be connected to the internet to use the agents module, even if you are using Ollama to run a locally-hosted LLM. TableMage's agent, ChatDA, relies on FastEmbed for retrieval augmented generation, but it may need to download the FastEmbed model from the internet prior to use. ChatDA can be run with a local LLM and FastEmbed, ensuring total data privacy.

Updates

TableMage is under active development.