ComputerAssignments/CA6_systematic_generalization/fix_notebook_errors.py

#!/usr/bin/env python3
"""
Fix specific errors in the CA6 notebook
"""

import nbformat
import re


def fix_notebook_errors(notebook_path):
    """Fix known errors in the notebook"""

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

    fixes_applied = []

    for i, cell in enumerate(nb.cells):
        if cell.cell_type == "code":
            original_source = cell.source

            # Fix 1: Cell with len(self.modules) error
            if (
                "len(self.modules)" in cell.source
                and "NeuralModuleNetwork" in cell.source
            ):
                cell.source = cell.source.replace(
                    "self.module_selector = nn.Linear(embed_dim, len(self.modules))",
                    "self.module_selector = nn.Linear(embed_dim, len(self.modules.keys()))",
                )
                if cell.source != original_source:
                    fixes_applied.append(f"Cell {i}: Fixed len(self.modules) issue")
                    print(
                        f"✅ Fixed Cell {i}: len(self.modules) -> len(self.modules.keys())"
                    )

            # Fix 2: Replace literal \n with actual newlines in large code blocks
            if "\\n" in cell.source and len(cell.source) > 500:
                # This is a cell with escaped newlines that should be actual newlines
                cell.source = cell.source.replace("\\n", "\n").replace("\\t", "\t")
                cell.source = cell.source.replace('\\"', '"').replace("\\'", "'")
                if cell.source != original_source:
                    fixes_applied.append(f"Cell {i}: Fixed escaped newlines")
                    print(
                        f"✅ Fixed Cell {i}: Converted escaped newlines to actual newlines"
                    )

            # Fix 3: Symbol evaluation issue - add better error handling
            if (
                "def evaluate(self, bindings" in cell.source
                and "class Compound" in cell.source
            ):
                # Add type checking before sum() operation
                if "return sum(eval_args)" in cell.source:
                    cell.source = cell.source.replace(
                        "        if self.operator == '+':\n            return sum(eval_args)",
                        """        if self.operator == '+':
            # Check if all args are numeric
            if all(isinstance(arg, (int, float)) for arg in eval_args):
                return sum(eval_args)
            else:
                return Compound(self.operator, *eval_args)""",
                    )
                    if cell.source != original_source:
                        fixes_applied.append(f"Cell {i}: Fixed Symbol addition issue")
                        print(
                            f"✅ Fixed Cell {i}: Added type checking for Symbol operations"
                        )

    if fixes_applied:
        print(f"\n{'='*80}")
        print(f"Applied {len(fixes_applied)} fixes:")
        for fix in fixes_applied:
            print(f"  • {fix}")

        # Save the fixed notebook
        with open(notebook_path, "w", encoding="utf-8") as f:
            nbformat.write(nb, f)
        print(f"\n✅ Saved fixed notebook to: {notebook_path}")
    else:
        print("No fixes needed or applied")

    return len(fixes_applied)


if __name__ == "__main__":
    notebook_path = "notebooks/CA6.ipynb"

    print("=" * 80)
    print("FIXING NOTEBOOK ERRORS")
    print("=" * 80)
    print()

    num_fixes = fix_notebook_errors(notebook_path)

    print(f"\n{'='*80}")
    print(f"Total fixes applied: {num_fixes}")
    print("=" * 80)