Unify solver and reorganize neural components

README.md

@@ -29,20 +29,23 @@ arc_solver_project/
 │   ├── dsl.py                 # Domain-specific language primitives
 │   ├── heuristics.py          # Heuristic rule inference
 │   ├── search.py              # Basic brute-force search
-│   ├── solver.py              # Main solver interface (enhanced)
-│   ├── enhanced_solver.py     # Enhanced solver with neural components
+│   ├── solver.py              # Main solver interface with enhancements
 │   ├── enhanced_search.py     # Neural-guided program synthesis
+│   ├── features.py            # Task feature extraction
+│   ├── ttt.py                 # Test-time training utilities
 │   ├── io_utils.py            # JSON loading and submission helpers
-│   └── neural/                # Neural guidance components
-│       ├── features.py        # Task feature extraction
+│   └── neural/                # Neural components
 │       ├── guidance.py        # Neural operation prediction
-│       ├── sketches.py        # Program sketch mining
 │       ├── episodic.py        # Episodic retrieval system
-│       └── ttt.py             # Test-time training
+│       └── sketches.py        # Program sketch mining
 │
 ├── arc_submit.py              # Command-line submission script
-├── train_neural_guidance.py   # Training script for neural components
-├── benchmark.py               # Benchmarking and evaluation tools
+├── tools/                     # Training and benchmarking utilities
+│   ├── train_guidance.py
+│   ├── mine_sketches.py
+│   ├── build_memory.py
+│   └── benchmark.py
+├── tests/                     # Unit and integration tests
 └── README.md                  # This file
 ```
 
@@ -62,22 +65,21 @@ ARC_USE_BASELINE=1 python arc_submit.py
 
 ```bash
 # Train neural guidance (requires training data)
-python train_neural_guidance.py
+python tools/train_guidance.py
 
 # Or setup environment with defaults
-python benchmark.py
+python tools/benchmark.py
 ```
 
 ### 3. Python API
 
 ```python
-from arc_solver.enhanced_solver import solve_task_enhanced
+from arc_solver.solver import solve_task_enhanced, ARCSolver
 
 # Solve a single task with full enhancements
 result = solve_task_enhanced(task)
 
 # Configure solver behavior
-from arc_solver.enhanced_solver import ARCSolver
 solver = ARCSolver(use_enhancements=True)
 result = solver.solve_task(task)
 ```

arc_solver/enhanced_search.py

@@ -15,9 +15,9 @@
 from .grid import Array, eq
 from .dsl import OPS, apply_program
 from .heuristics import consistent_program_single_step, score_candidate, diversify_programs
-from .guidance import NeuralGuidance
-from .memory import EpisodicRetrieval
-from .sketches import SketchMiner, generate_parameter_grid
+from .neural.guidance import NeuralGuidance
+from .neural.episodic import EpisodicRetrieval
+from .neural.sketches import SketchMiner, generate_parameter_grid
 from .ttt import TestTimeTrainer, DataAugmentation
 
 

arc_solver/neural/__init__.py

@@ -0,0 +1,17 @@
+"""Neural components for the ARC solver."""
+
+from .guidance import SimpleClassifier, HeuristicGuidance, NeuralGuidance
+from .episodic import Episode, EpisodeDatabase, EpisodicRetrieval
+from .sketches import ProgramSketch, SketchMiner, generate_parameter_grid
+
+__all__ = [
+    "SimpleClassifier",
+    "HeuristicGuidance",
+    "NeuralGuidance",
+    "Episode",
+    "EpisodeDatabase",
+    "EpisodicRetrieval",
+    "ProgramSketch",
+    "SketchMiner",
+    "generate_parameter_grid",
+]

arc_solver/memory.py → arc_solver/neural/episodic.py

@@ -14,15 +14,15 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Tuple, Optional
+from typing import Any, Dict, List, Optional, Tuple
 from collections import defaultdict
 import json
 import os
 
 import numpy as np
 
-from .grid import Array
-from .features import compute_task_signature, extract_task_features
+from ..grid import Array
+from ..features import compute_task_signature, extract_task_features
 
 # Type alias for a DSL program representation used across the project
 Program = List[Tuple[str, Dict[str, Any]]]