working.md

BiasBalance - How Everything Works Together 🔍

Overview

BiasBalance is a comprehensive Streamlit-based web application designed to detect bias in categorical datasets and provide AI-powered data augmentation solutions. The application follows a modular architecture with clear separation of concerns across multiple Python modules.

🏗️ Architecture Overview

BiasBalance/
├── src/
│   ├── Home.py                 # Main entry point & dataset management
│   ├── components.py           # Reusable UI components
│   ├── utils.py               # Core statistical functions & utilities
│   ├── data_augmentation.py   # AI-powered data generation
│   └── pages/
│       ├── 1_Dataset_Info.py      # Dataset exploration & column selection
│       ├── 2_Goodness_of_Fit_Test.py  # Bias detection using Chi-square
│       ├── 3_Independence_Test.py     # Variable relationship analysis
│       └── 4_Final_Augmentation.py   # AI data generation interface
├── requirements.txt           # Python dependencies
├── Dockerfile                # Container configuration
├── docker-compose.yml       # Multi-container orchestration
└── test-dataset.csv         # Sample dataset for testing

🔄 Application Flow

1. Entry Point & Data Loading (Home.py)

The application starts with Home.py, which serves as both the main entry point and the central data management hub.

Key Responsibilities:

• Session State Initialization: Sets up global application state using Streamlit's session management
• Dataset Upload: Handles CSV file uploads and validates data format
• Test Dataset Loading: Provides a pre-loaded sample dataset for demonstration
• Data Preprocessing: Automatically detects and categorizes columns based on data types and cardinality
• Download Management: Enables users to download augmented datasets

Flow:

1. User uploads CSV or loads test dataset
2. System analyzes columns and identifies categorical vs numerical data
3. Default categorical columns are auto-selected based on heuristics
4. Dataset is stored in session state for cross-page access

2. Dataset Exploration (1_Dataset_Info.py)

This page provides comprehensive dataset analysis and column configuration.

Key Features:

• Statistical Overview: Displays dataset dimensions, column types, and data quality metrics
• Column Selection Interface: Multi-select widget for choosing which columns to analyze
• Data Quality Assessment: Shows null value percentages and unique value counts
• Preview Functionality: Displays data samples for user verification

Selection Criteria:

• Default Selected: Categorical/Object columns with >50% non-null values
• Optional: Other categorical columns + numeric columns with <25% unique values

3. Bias Detection (2_Goodness_of_Fit_Test.py)

The core statistical analysis module that detects bias using Chi-square Goodness-of-Fit tests.

Statistical Foundation:

• Null Hypothesis: Data follows the expected distribution (equal or custom proportions)
• Test Statistic: Chi-square = Σ((Observed - Expected)² / Expected)
• Residual Analysis: Standardized residuals identify over/under-represented categories

Workflow:

1. Column Selection: User chooses which categorical column to analyze
2. Distribution Configuration:
   • Equal distribution (uniform expected frequencies)
   • Proportional distribution (custom expected percentages)
3. Statistical Testing: Calculates chi-square statistic and p-value
4. Result Interpretation:
   • Significant bias detection (p < α)
   • Category-level analysis using residuals
   • Augmentation recommendations for under-represented categories

Residual Thresholds:

• Positive Threshold (+1.3): Over-represented categories
• Negative Threshold (-1.3): Under-represented categories (candidates for augmentation)

4. Independence Testing (3_Independence_Test.py)

Analyzes relationships between pairs of categorical variables using Chi-square tests of independence.

Statistical Approach:

• Contingency Table: Cross-tabulation of two categorical variables
• Expected Frequencies: Calculated under independence assumption
• Test Statistic: Chi-square test of independence
• Association Analysis: Standardized residuals reveal specific associations

Output:

• Statistical significance of relationship
• Detailed residual analysis showing specific category associations
• Visual representation through contingency tables

5. AI-Powered Data Augmentation (4_Final_Augmentation.py)

The most sophisticated component that generates synthetic data using OpenAI's GPT models.

Architecture:

User Input → Prompt Generation → LLM API → Pydantic Validation → DataFrame Output

Key Components:

Augmentation Needs Calculation (data_augmentation.py)

def calculate_augmentation_needs(residuals, observed, expected, categories, threshold):
    # For each under-represented category (residual < -1.3):
    # Calculate: samples_needed = ceil((expected - observed) * 1.1)

Multi-Column Optimization

The system intelligently handles multiple biased columns by:

1. Combining Requirements: Merges augmentation needs across all saved test results
2. Optimization: Generates minimum samples that satisfy all column requirements simultaneously
3. Distribution Preservation: Ensures new data maintains realistic relationships between variables

Pydantic Model Generation

Dynamic model creation ensures type safety and validation:

def create_pydantic_model_from_df(df, category_info, combined_needs):
    # Creates Enum types for categorical columns
    # Validates generated data meets count requirements
    # Ensures data integrity through model validation

LLM Integration

• Model Selection: Supports both GPT-4o (complex) and GPT-4o-mini (fast) models
• Structured Output: Uses OpenAI's function calling for guaranteed JSON format
• Error Handling: Implements retry logic with improved prompts on validation failures
• API Key Management: Secure session-based key storage

🧩 Core Modules Deep Dive

utils.py - Statistical Engine

The mathematical foundation of the application:

Bias Detection Algorithm:

def detect_bias(data, categories, distribution, proportions):
    # 1. Data preprocessing (handle nulls, type conversion)
    # 2. Calculate observed frequencies
    # 3. Determine expected frequencies (equal vs proportional)
    # 4. Perform chi-square test using scipy.stats.chisquare
    # 5. Calculate standardized residuals
    return chi2_stat, p_value, residuals, observed, expected

Independence Testing:

def test_independence(data, col1, col2):
    # 1. Create contingency table
    # 2. Apply chi-square test of independence
    # 3. Calculate expected frequencies under independence
    # 4. Compute standardized residuals
    return chi2, p_value, residuals, contingency_table, expected

components.py - UI Framework

Provides consistent user interface elements:

• Sidebar Management: Persistent navigation and settings across all pages
• Analysis Settings: Configurable significance levels and residual thresholds
• Dataset Status: Real-time feedback on loaded data and analysis state

data_augmentation.py - AI Engine

The most complex module handling synthetic data generation:

Prompt Engineering:

system_message = """Create synthetic data to oversample under-represented 
categories... minimize redundant samples... respect relationships between columns..."""

Data Generation Pipeline:

1. Needs Analysis: Calculate augmentation requirements per category
2. Prompt Construction: Build comprehensive prompts with context and constraints
3. LLM Invocation: Call OpenAI API with structured output format
4. Validation: Ensure generated data meets all requirements using Pydantic
5. Integration: Combine synthetic data with original dataset

🔄 Session State Management

Streamlit's session state enables cross-page data persistence:

# Core data storage
st.session_state.df                    # Main dataset
st.session_state.categorical_cols      # Selected categorical columns
st.session_state.saved_test_results   # Bias test results for augmentation

# Analysis parameters
st.session_state.alpha                 # Significance level
st.session_state.positive_threshold    # Over-representation threshold
st.session_state.negative_threshold    # Under-representation threshold

# Augmentation state
st.session_state.generated_response    # AI-generated synthetic data
st.session_state.openai_api_key       # Secure API key storage

🔬 Statistical Methodology

Chi-Square Goodness-of-Fit Test

• Purpose: Detect deviations from expected categorical distributions
• Assumptions: Independent observations, adequate sample size (expected ≥ 5)
• Interpretation: p < α indicates significant bias

Standardized Residuals

residual = (observed - expected) / sqrt(expected)

• Interpretation: |residual| > 2 indicates significant deviation
• Application: Identifies specific categories requiring augmentation

Chi-Square Test of Independence

• Purpose: Detect associations between categorical variables
• H₀: Variables are independent
• H₁: Variables are associated

🚀 Deployment & Scalability

Docker Configuration

FROM python:3.11-slim
# Optimized for production deployment
# Supports multi-architecture builds (linux/amd64, linux/arm64)

Environment Management

• Development: Local Python virtual environment
• Production: Docker container with volume mounting
• CI/CD: Automated GitHub Actions for Docker Hub publishing

🔧 Data Flow Visualization

CSV Upload → Column Analysis → Categorical Selection
     ↓
Bias Testing → Statistical Analysis → Result Storage
     ↓
Multiple Tests → Combined Requirements → AI Prompt Generation
     ↓
LLM Processing → Data Validation → Dataset Augmentation
     ↓
Download/Export → Enhanced Dataset

🎯 Key Design Decisions

1. Modular Architecture: Separation of statistical logic, UI components, and AI integration
2. Session-Based State: Enables complex multi-page workflows without database requirements
3. Dynamic Model Generation: Pydantic models ensure type safety for diverse datasets
4. Intelligent Optimization: Multi-column augmentation minimizes synthetic data requirements
5. Error Recovery: Robust retry mechanisms for LLM integration
6. User-Centric Design: Progressive disclosure of complexity through intuitive UI

🔮 Extension Points

The architecture supports easy extension for:

• Additional Statistical Tests: New test modules can be added to the pages/ directory
• Alternative LLM Providers: The AI interface is abstracted for provider swapping
• Export Formats: Additional data export options (JSON, Parquet, etc.)
• Visualization: Enhanced charts and statistical plots
• Batch Processing: Multiple dataset analysis workflows

This design creates a powerful, user-friendly tool that bridges the gap between statistical analysis and modern AI capabilities for dataset enhancement and bias correction.