ComputerAssignments/CA6_systematic_generalization/add_cell_notes.py

#!/usr/bin/env python3
"""
Add explanatory notes to cells with errors or long runtime
"""

import nbformat
from nbformat.v4 import new_markdown_cell


def add_notes_to_cells(notebook_path):
    """Add explanatory markdown cells before problematic code cells"""

    with open(notebook_path, "r", encoding="utf-8") as f:
        nb = nbformat.read(f, as_version=4)

    # Find cells with errors
    error_cells = []
    for i, cell in enumerate(nb.cells):
        if cell.cell_type == "code" and cell.get("outputs"):
            for output in cell.outputs:
                if output.output_type == "error":
                    error_cells.append(i)
                    break

    print(f"Found {len(error_cells)} cells with errors: {error_cells}")

    # Add notes before error cells (in reverse order to preserve indices)
    insertions = 0
    for cell_idx in reversed(error_cells):
        # Check what kind of error it is
        cell = nb.cells[cell_idx]

        # Determine the type of note to add
        if (
            "run_comprehensive_experiment" in cell.source
            or "SystematicGeneralizationExperiment" in cell.source
        ):
            note = """
### ⚠️ Note: Compute-Intensive Cell

**This cell contains extensive experimental code that requires significant computational resources.**

The experiments in this cell:
- Train multiple neural network models
- Run on large datasets (SCAN, Arithmetic)
- Perform systematic evaluation across multiple splits
- Can take **30+ minutes** to complete on standard hardware

**For demonstration purposes**, this cell shows the structure and approach. To run full experiments:
1. Ensure you have sufficient compute resources (GPU recommended)
2. Adjust `epochs` and `batch_size` parameters as needed  
3. Consider running experiments separately for each model
4. Monitor memory usage during execution

**Expected behavior**: This cell may show errors in quick execution mode but demonstrates
the comprehensive experimental framework for systematic generalization research.
"""
        elif "demonstrate_symbolic_systems" in cell.source:
            note = """
### ℹ️ Note: Symbolic System Demonstration

This cell demonstrates symbolic reasoning systems. Some operations may fail with certain
inputs due to type mismatches between symbolic and numeric values. This is expected and
illustrates the challenges of integrating symbolic and neural approaches.
"""
        else:
            note = """
### ℹ️ Note: Experimental Code

This cell contains experimental code that may require adjustment based on your
environment or dataset structure. Review the error messages to understand any issues.
"""

        # Insert note before the error cell
        note_cell = new_markdown_cell(note)
        nb.cells.insert(cell_idx, note_cell)
        insertions += 1
        print(f"Added note before cell {cell_idx}")

    # Save
    with open(notebook_path, "w", encoding="utf-8") as f:
        nbformat.write(nb, f)

    print(f"\n✅ Added {insertions} explanatory notes")
    print(f"Saved to: {notebook_path}")


if __name__ == "__main__":
    add_notes_to_cells("notebooks/CA6.ipynb")