REFACTORING_PLAN.md

Code Organization and Refactoring Plan

This document outlines the plan for reorganizing the codebase to improve structure, maintainability, and testability, with a focus on resolving issues found during testing.

Current Issues

1. Code Organization Issues

   • /sentinel package contains a partial implementation of a cleaner architecture
   • Many modules are scattered across the codebase without a clear organization
   • New neural plasticity functionality needs proper integration

2. Circular Import Problems

   • Several circular imports that cause errors, especially with the datasets package
   • Import path management needs improvement
   • Inconsistent use of relative vs. absolute imports

3. Inconsistent API Design

   • Mixed responsibilities in modules
   • Redundant implementations (fine_tuner vs. fine_tuner_improved)
   • Inconsistent function signatures and parameter names

4. Testing Limitations

   • Lack of robust test infrastructure
   • No CI/CD pipeline
   • Manual testing procedures prone to error

Target Structure

sentinel/
├── __init__.py
├── models/
│   ├── __init__.py
│   ├── adaptive/
│   ├── loaders/
│   └── optimized/
├── controller/
│   ├── __init__.py
│   ├── metrics/
│   └── visualizations/
├── pruning/
│   ├── __init__.py
│   ├── fixed_pruning_module.py
│   ├── fixed_pruning_module_jax.py
│   ├── pruning_module.py
│   ├── strategies/
│   │   ├── __init__.py
│   │   ├── base.py
│   │   ├── entropy.py
│   │   ├── magnitude.py
│   │   └── random.py
│   ├── growth/
│   │   ├── __init__.py
│   │   ├── base.py
│   │   ├── utils.py
│   │   ├── gradient_sensitivity.py
│   │   ├── entropy_gap.py
│   │   ├── balanced.py
│   │   └── random.py
│   ├── fine_tuning/
│   │   ├── __init__.py
│   │   └── fine_tuner.py  # Consolidated implementation
│   ├── stability/
│   │   ├── __init__.py
│   │   ├── memory_management.py
│   │   └── nan_prevention.py
│   └── visualization/
│       ├── __init__.py
│       └── visualization.py
├── data/
│   ├── __init__.py
│   └── loaders/
│       ├── __init__.py
│       └── dataset_utils.py  # Robust dataset loading without circular imports
└── utils/
    ├── __init__.py
    ├── metrics.py
    ├── progress_tracker.py
    └── checkpoints/

Known Issues and Solutions

1. Datasets Import Error:

   • Error: Cannot import datasets: name 'datasets' is not defined
   • Solution: Implemented a mock datasets module to break circular imports in sentinel/upgrayedd/utils/data.py, which should be moved to sentinel/data/loaders/dataset_utils.py

2. Tokenizer Padding Issue:

   • Error: ValueError: Asking to pad but the tokenizer does not have a padding token
   • Solution: Set tokenizer.pad_token = tokenizer.eos_token for GPT-2 models in a utility function

3. Attention Mask Warning:

   • Warning: The attention mask is not set and cannot be inferred from input because pad token is same as eos token
   • Solution: Explicitly create attention masks when pad_token equals eos_token

4. Leaked Semaphore Objects:

   • Warning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
   • Solution: This is a minor issue with torch multiprocessing that can be ignored

Migration Plan

Phase 1: Fix Circular Imports and Dataset Loading

1. Create robust dataset loading utilities:

   • Create sentinel/data/loaders/dataset_utils.py with the improved mock datasets approach
   • Add utility functions for preparing data with proper attention masks
   • Ensure backward compatibility with existing code

2. Update import patterns in all modules:

   • Use absolute imports for standard libraries and third-party packages
   • Use relative imports for internal modules where appropriate
   • Avoid importing from . when possible to reduce confusion

Phase 2: Consolidate Fine Tuner Implementations

1. Create a single robust fine tuner:

   • Merge functionality from fine_tuner.py, fine_tuner_improved.py, and fine_tuner_consolidated.py
   • Add a stability_level parameter to control robustness features
   • Ensure all edge cases are handled properly

2. Add comprehensive logging and error handling:

   • Improve error messages for common issues
   • Add more detailed logging for debugging
   • Implement graceful fallbacks for non-critical errors

Phase 3: Robust Testing Framework

1. Create unit tests for core components:

   • Test strategies module (entropy, magnitude, random)
   • Test data loading with mocks
   • Test fine tuning with small models

2. Create integration tests:

   • End-to-end tests with minimal example models
   • Test full pipeline from loading to pruning to fine-tuning
   • Add performance regression tests

3. Set up continuous integration:

   • Configure automated tests on push
   • Add linting and code quality checks
   • Create test coverage reports

Phase 4: Documentation and Examples

1. Update documentation:

   • Create clear API documentation for the main modules
   • Add examples for common use cases
   • Update README with new structure and usage information

2. Create tutorial notebooks:

   • Simple examples with different models
   • Advanced use cases and customization
   • Performance optimization guides

Testing Commands

For testing during the migration, use these commands:

# Test dataset loading
python test_dataset_loading.py

# Test experiment runner
python test_run_experiment.py

# Test notebook integration
python colab_notebooks/PruningAndFineTuningColab.py --test_mode --super_simple --model_name distilgpt2

# Run full experiment (when fixed)
python sentinel/scripts/run_experiment.py --model_name distilgpt2 --pruning_ratio 0.2 --strategy random --max_cycles 1 --epochs 1 --device cpu

Backward Compatibility

To ensure backward compatibility, we'll:

1. Keep old module locations but have them import from new locations
2. Maintain the same class and function signatures
3. Add deprecation warnings to old locations indicating the recommended import path

Testing Strategy

Each phase will be tested to ensure:

1. All existing functionality works as before
2. All scripts can run without modifications
3. The test suite passes with the same results

🤖 Generated with Claude Code