Model

Model abstraction, metadata tracking, and serialization for sklearn, LightGBM, XGBoost, and other frameworks. ModelSerializer handles save/load with full metadata and optional dataset storage. ModelRegistry tracks model lineage across runs.

Base Model

bitbullet.model.core.base.BaseModel

Bases: ABC

Abstract base class for all model wrappers.

Provides a consistent interface across different ML frameworks, enabling seamless switching between sklearn, LightGBM, XGBoost, R, etc.

Key features: - Unified predict/predict_proba API - Feature name tracking and validation - Serialization support - Metadata tracking

Example

from bitbullet.model.wrappers.lgbm_wrapper import LGBMClassifierWrapper

model = LGBMClassifierWrapper(n_estimators=100, max_depth=5)
model.fit(X_train, y_train)
predictions = model.predict(X_test)

classes_ property

Get the class labels (for classifiers).

feature_names property

Get the feature names the model was trained on.

framework property

Get the framework name (e.g., 'lightgbm').

hyperparameters property

Get the model hyperparameters.

is_fitted property

Check if model has been fitted.

model_type property

Get the model type (e.g., 'LGBMClassifier').

n_features property

Get the number of features the model expects.

state property

Get the complete model state.

__init__(**kwargs)

Initialize model wrapper.

Parameters:

Name	Type	Description	Default
**kwargs		Model-specific parameters	{}

__repr__()

String representation of the model.

fit(X, y, **kwargs) abstractmethod

Fit the model to training data.

Parameters:

Name	Type	Description	Default
X	Union[DataFrame, ndarray]	Training features	required
y	Union[Series, ndarray]	Training targets	required
**kwargs		Additional fitting parameters (sample_weight, eval_set, etc.)	{}

Returns:

Name	Type	Description
self	BaseModel	Fitted model instance

get_params()

Get model parameters (sklearn-compatible).

Returns:

Type	Description
dict	Dictionary of model parameters

predict(X) abstractmethod

Generate predictions for input data.

Parameters:

Name	Type	Description	Default
X	Union[DataFrame, ndarray]	Input features	required

Returns:

Type	Description
ndarray	Predicted labels

predict_proba(X) abstractmethod

Generate probability predictions for input data.

Parameters:

Name	Type	Description	Default
X	Union[DataFrame, ndarray]	Input features	required

Returns:

Type	Description
ndarray	Class probabilities (n_samples, n_classes)

set_params(**params)

Set model parameters (sklearn-compatible).

Parameters:

Name	Type	Description	Default
**params		Parameters to set	{}

Returns:

Name	Type	Description
self	BaseModel	Model instance

Metadata

bitbullet.model.core.metadata.ModelMetadata dataclass

Comprehensive metadata for a trained model.

Tracks everything needed for reproducibility, auditing, and model management in production environments.

Attributes:

Name	Type	Description
name	str	Model name/identifier
version	str	Model version
model_type	str	Type of model (e.g., 'LGBMClassifier')
framework	str	Framework used (e.g., 'lightgbm')
task	str	Task type ('classification' or 'regression')
created_at	datetime	When the model was created
trained_at	Optional[datetime]	When the model was last trained
hyperparameters	Dict[str, Any]	Model hyperparameters
feature_names	List[str]	List of features used
n_features	int	Number of features
classes	Optional[List[Any]]	Class labels (for classification)
metrics	Dict[str, float]	Performance metrics
training_time_seconds	float	Time taken to train
cv_folds	int	Number of CV folds used
cv_scores	Dict[str, List[float]]	Cross-validation scores
optimal_threshold	Optional[float]	Optimal classification threshold
training_dataset	Optional[DatasetMetadata]	Metadata for training data
validation_dataset	Optional[DatasetMetadata]	Metadata for validation data
test_dataset	Optional[DatasetMetadata]	Metadata for test data
environment	Dict[str, str]	Environment information (Python version, packages)
tags	List[str]	Custom tags for organization
notes	str	Free-form notes

add_artifact_metadata(**metadata)

Merge arbitrary artifact metadata into the model metadata.

add_cv_scores(metric_name, scores)

Add cross-validation scores.

Parameters:

Name	Type	Description	Default
metric_name	str	Name of the metric	required
scores	List[float]	List of scores from each fold	required

add_dataset(X, y=None, dataset_type='train')

Add dataset metadata.

Parameters:

Name	Type	Description	Default
X	DataFrame	Feature DataFrame	required
y	Optional[Series]	Target Series (optional)	None
dataset_type	str	One of 'train', 'validation', 'test'	'train'

add_feature_schema(X)

Store a JSON-friendly feature schema from a feature frame.

add_metric(name, value)

Add a performance metric.

Parameters:

Name	Type	Description	Default
name	str	Metric name (e.g., 'roc_auc', 'accuracy')	required
value	float	Metric value	required

summary()

Generate a human-readable summary of the model.

Returns:

Type	Description
str	Formatted summary string

to_dict()

Convert metadata to dictionary.

bitbullet.model.core.metadata.DatasetMetadata dataclass

Metadata for a dataset (train/validation/test).

Tracks essential information about datasets used in model training and evaluation, enabling full reproducibility and audit trails.

Attributes:

Name	Type	Description
name	str	Dataset identifier (e.g., 'train', 'validation', 'test')
n_samples	int	Number of samples
n_features	int	Number of features
feature_names	List[str]	List of feature names
feature_dtypes	Dict[str, str]	Data types of features
target_name	Optional[str]	Name of the target variable
target_dtype	Optional[str]	Data type of the target
class_distribution	Dict[Any, int]	Distribution of classes (for classification)
date_range	Optional[tuple]	Date range of the data (if temporal)
created_at	datetime	When this metadata was created
data_hash	Optional[str]	Hash of the data for integrity checking
statistics	Dict[str, Any]	Summary statistics of features

from_dataframe(X, y=None, name='unknown', compute_hash=True) classmethod

Create metadata from a pandas DataFrame.

Parameters:

Name	Type	Description	Default
X	DataFrame	Feature DataFrame	required
y	Optional[Series]	Target Series (optional)	None
name	str	Dataset name	'unknown'
compute_hash	bool	Whether to compute data hash	True

Returns:

Type	Description
DatasetMetadata	DatasetMetadata instance

to_dict()

Convert metadata to dictionary.

Registry

bitbullet.model.core.registry.ModelRegistry

Registry for model wrappers.

Allows users to register custom model wrappers and create them dynamically by name.

Example

from bitbullet.model.core.registry import ModelRegistry

# Register a custom model
@ModelRegistry.register("my_custom_model")
class MyCustomModel(BaseModel):
    def fit(self, X, y, **kwargs):
        # Custom training logic
        pass

    def predict(self, X):
        # Custom prediction logic
        pass

    def predict_proba(self, X):
        # Custom probability prediction
        pass

# Create an instance
model = ModelRegistry.create("my_custom_model", param1=value1)

create(name, **kwargs) classmethod

Create a model instance by name.

Parameters:

Name	Type	Description	Default
name	str	Registered model name	required
**kwargs		Model parameters	{}

Returns:

Type	Description
BaseModel	Model instance

Raises:

Type	Description
ValueError	If model name not found

is_registered(name) classmethod

Check if a model is registered.

Parameters:

Name	Type	Description	Default
name	str	Model name to check	required

Returns:

Type	Description
bool	True if registered, False otherwise

list_models() classmethod

List all registered models.

Returns:

Type	Description
Dict[str, Type[BaseModel]]	Dictionary mapping model names to classes

register(name) classmethod

Decorator to register a model wrapper.

Parameters:

Name	Type	Description	Default
name	str	Name to register the model under	required

Returns:

Type	Description
	Decorator function

Example

@ModelRegistry.register("lgbm_classifier")
class LGBMClassifierWrapper(BaseModel):
    pass

unregister(name) classmethod

Remove a model from the registry.

Parameters:

Name	Type	Description	Default
name	str	Model name to unregister	required

Raises:

Type	Description
ValueError	If model not found

Serialization

bitbullet.model.serialization.model_serializer.ModelSerializer

Handles saving and loading of models with complete metadata.

Supports multiple serialization formats (pickle, joblib) and can optionally include training/validation/test datasets for complete reproducibility.

Example

from bitbullet.model.serialization import ModelSerializer

# Save model with all datasets
ModelSerializer.save(
    model=trained_model,
    path="./models/fraud_model_v1.pkl",
    metadata=metadata,
    train_data=(X_train, y_train),
    test_data=(X_test, y_test),
    include_datasets=True
)

# Load model package
package = ModelSerializer.load("./models/fraud_model_v1.pkl")
print(package.summary())

# Access components
model = package.model
test_X = package.test_data['X']
test_y = package.test_data['y']

load(path) staticmethod

Load a saved model package.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Path to saved model	required

Returns:

Type	Description
ModelPackage	ModelPackage containing model, metadata, and datasets

Raises:

Type	Description
FileNotFoundError	If model file not found

Example

package = ModelSerializer.load("./models/model_v1.pkl")

# Access model
model = package.model
predictions = model.predict(X_new)

# Access metadata
print(package.metadata.summary())

# Access datasets (if included)
if package.test_data:
    X_test = package.test_data['X']
    y_test = package.test_data['y']

load_metadata(path) staticmethod

Load only the metadata from a saved model.

Useful for inspecting model details without loading the full model.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Path to saved model or metadata JSON	required

Returns:

Type	Description
ModelMetadata	ModelMetadata instance

load_model_only(path) staticmethod

Load a model-only artifact.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Path to model file	required

Returns:

Type	Description
BaseModel	BaseModel instance

save(model, path, metadata=None, train_data=None, validation_data=None, test_data=None, include_datasets=True, format='pickle', compress=True) staticmethod

Save model with comprehensive metadata and optional datasets.

Parameters:

Name	Type	Description	Default
model	BaseModel	Fitted model to save	required
path	Union[str, Path]	Save path (will create parent directories)	required
metadata	Optional[ModelMetadata]	Model metadata (will create basic if None)	None
train_data	Optional[tuple]	Tuple of (X_train, y_train) to save with model	None
validation_data	Optional[tuple]	Tuple of (X_val, y_val) to save	None
test_data	Optional[tuple]	Tuple of (X_test, y_test) to save	None
include_datasets	bool	Whether to include datasets in saved file	True
format	str	Serialization format ('pickle' or 'joblib')	'pickle'
compress	bool	Whether to compress (joblib only)	True

Example

ModelSerializer.save(
    model=lgbm_model,
    path="./models/model_v1.pkl",
    metadata=metadata,
    train_data=(X_train, y_train),
    test_data=(X_test, y_test),
    include_datasets=True
)

save_model_only(model, path, format='pickle') staticmethod

Save only the model (no metadata or datasets).

Useful for lightweight deployments where metadata isn't needed.

Parameters:

Name	Type	Description	Default
model	BaseModel	Model to save	required
path	Union[str, Path]	Save path	required
format	str	'pickle' or 'joblib'	'pickle'

Wrappers

!!! note Wrappers are only available when the corresponding optional extra is installed. Install with pip install "bitbullet[inference-models]".

bitbullet.model.wrappers.lgbm_wrapper.LGBMClassifierWrapper

Bases: BaseModel

Wrapper for LightGBM Classifier with enhanced functionality.

Provides consistent API with feature tracking, validation, and metadata management.

Example

from bitbullet.model.wrappers.lgbm_wrapper import LGBMClassifierWrapper

model = LGBMClassifierWrapper(
    n_estimators=100,
    max_depth=5,
    learning_rate=0.1
)

model.fit(X_train, y_train,
          eval_set=[(X_val, y_val)],
          callbacks=[lgb.early_stopping(20)])

predictions = model.predict(X_test)
probabilities = model.predict_proba(X_test)

best_iteration_ property

Get the best iteration from early stopping.

Returns:

Type	Description
Optional[int]	Best iteration number or None

best_score_ property

Get the best score from early stopping.

Returns:

Type	Description
Optional[dict]	Best score dictionary or None

booster_ property

Get the underlying LightGBM Booster.

Returns:

Type	Description
	LightGBM Booster object

feature_importances_ property

Get feature importances (split-based).

Returns:

Type	Description
ndarray	Feature importance scores

__init__(**kwargs)

Initialize LightGBM classifier wrapper.

Parameters:

Name	Type	Description	Default
**kwargs		LightGBM parameters (n_estimators, max_depth, etc.)	{}

fit(X, y, **kwargs)

Fit the LightGBM classifier.

Parameters:

Name	Type	Description	Default
X	Union[DataFrame, ndarray]	Training features	required
y	Union[Series, ndarray]	Training targets	required
**kwargs		Additional fitting parameters: - sample_weight: Sample weights - eval_set: List of (X, y) tuples for validation - callbacks: List of callback functions - init_model: LightGBM booster to continue training from	{}

Returns:

Name	Type	Description
self	LGBMClassifierWrapper	Fitted model instance

get_feature_importance(importance_type='split')

Get feature importance as DataFrame.

Parameters:

Name	Type	Description	Default
importance_type	str	'split' or 'gain'	'split'

Returns:

Type	Description
DataFrame	DataFrame with features and their importance scores

predict(X)

Generate predictions.

Parameters:

Name	Type	Description	Default
X	Union[DataFrame, ndarray]	Input features	required

Returns:

Type	Description
ndarray	Predicted class labels

Raises:

Type	Description
RuntimeError	If model not fitted
ValueError	If feature mismatch

predict_proba(X)

Generate probability predictions.

Parameters:

Name	Type	Description	Default
X	Union[DataFrame, ndarray]	Input features	required

Returns:

Type	Description
ndarray	Class probabilities (n_samples, n_classes)

Raises:

Type	Description
RuntimeError	If model not fitted
ValueError	If feature mismatch

bitbullet.model.wrappers.lgbm_wrapper.LGBMRegressorWrapper

Bases: BaseModel

Wrapper for LightGBM Regressor.

Similar to LGBMClassifierWrapper but for regression tasks.

best_iteration_ property

Get the best iteration from early stopping.

best_score_ property

Get the best score from early stopping.

feature_importances_ property

Get feature importances.

__init__(**kwargs)

Initialize LightGBM regressor wrapper.

Parameters:

Name	Type	Description	Default
**kwargs		LightGBM parameters	{}

fit(X, y, **kwargs)

Fit the LightGBM regressor.

get_feature_importance(importance_type='split')

Get feature importance as DataFrame.

predict(X)

Generate predictions.

predict_proba(X)

Not applicable for regression. Raises NotImplementedError.

bitbullet.model.wrappers.xgb_wrapper.XGBClassifierWrapper

Bases: BaseModel

Wrapper for XGBoost Classifier with enhanced functionality.

Example

from bitbullet.model.wrappers.xgb_wrapper import XGBClassifierWrapper

model = XGBClassifierWrapper(
    n_estimators=100,
    max_depth=5,
    learning_rate=0.1
)

model.fit(X_train, y_train,
          eval_set=[(X_val, y_val)],
          early_stopping_rounds=20)

predictions = model.predict(X_test)

best_iteration property

Get the best iteration from early stopping.

feature_importances_ property

Get feature importances.

__init__(**kwargs)

Initialize XGBoost classifier wrapper.

Parameters:

Name	Type	Description	Default
**kwargs		XGBoost parameters	{}

fit(X, y, **kwargs)

Fit the XGBoost classifier.

Parameters:

Name	Type	Description	Default
X	Union[DataFrame, ndarray]	Training features	required
y	Union[Series, ndarray]	Training targets	required
**kwargs		Additional fitting parameters	{}

Returns:

Name	Type	Description
self	XGBClassifierWrapper	Fitted model instance

predict(X)

Generate predictions.

predict_proba(X)

Generate probability predictions.

bitbullet.model.wrappers.xgb_wrapper.XGBRegressorWrapper

Bases: BaseModel

Wrapper for XGBoost Regressor.

best_iteration property

Get the best iteration from early stopping.

feature_importances_ property

Get feature importances.

__init__(**kwargs)

Initialize XGBoost regressor wrapper.

fit(X, y, **kwargs)

Fit the XGBoost regressor.

predict(X)

Generate predictions.

predict_proba(X)

Not applicable for regression.