demo.py

#!/usr/bin/env python3
"""
Quick demo of the HRM implementation.
"""

import torch
import sys
from pathlib import Path

# Add src to path
sys.path.append(str(Path(__file__).parent / 'src'))

from hrm.models.hrm import HierarchicalReasoningModel
from hrm.tasks.sudoku import SudokuAdapter, SudokuDataset
from hrm.tasks.maze import MazeAdapter, MazeDataset

def demo_sudoku():
    """Demo Sudoku solving"""
    print("="*60)
    print("SUDOKU SOLVING DEMO")
    print("="*60)
    
    # Create model and adapter
    model = HierarchicalReasoningModel(
        input_dim=512,
        output_dim=512,
        high_hidden_dim=512,
        low_hidden_dim=256,
        num_iterations=3
    )
    
    adapter = SudokuAdapter(
        hrm_input_dim=512,
        hrm_output_dim=512
    )
    
    # Create a small dataset
    dataset = SudokuDataset(num_puzzles=10, difficulty='easy')
    
    # Get a sample puzzle
    sample = dataset[0]
    puzzle = sample['puzzle'].unsqueeze(0)  # Add batch dimension
    solution = sample['solution']
    
    print(f"Input puzzle shape: {puzzle.shape}")
    print(f"Puzzle (partial):")
    print(puzzle[0, :3, :3])  # Show top-left 3x3 block
    
    # Encode puzzle
    encoded = adapter.encode_puzzle(puzzle)
    print(f"\nEncoded shape: {encoded.shape}")
    
    # Forward pass through HRM
    with torch.no_grad():
        output = model(encoded)
    
    print(f"HRM output shape: {output['output'].shape}")
    
    # Decode solution
    predicted, logits = adapter.decode_solution(
        output['output'],
        puzzle,
        enforce_constraints=True
    )
    
    print(f"\nPredicted solution shape: {predicted.shape}")
    print(f"Predicted (partial):")
    print(predicted[0, :3, :3])  # Show top-left 3x3 block
    
    # Check accuracy (untrained model, so random)
    correct = (predicted[0] == solution).float().mean()
    print(f"\nAccuracy (untrained): {correct:.2%}")
    
    return model, adapter

def demo_maze():
    """Demo Maze solving"""
    print("\n" + "="*60)
    print("MAZE SOLVING DEMO")
    print("="*60)
    
    # Create model and adapter
    model = HierarchicalReasoningModel(
        input_dim=512,
        output_dim=512,
        high_hidden_dim=512,
        low_hidden_dim=256,
        num_iterations=3
    )
    
    adapter = MazeAdapter(
        max_size=50,
        hrm_input_dim=512,
        hrm_output_dim=512
    )
    
    # Create a small maze dataset
    dataset = MazeDataset(num_mazes=10, min_size=10, max_size=20)
    
    # Get a sample maze
    sample = dataset[0]
    maze = sample['maze'].unsqueeze(0)  # Add batch dimension
    solution_path = sample['solution']
    
    print(f"Maze shape: {maze.shape}")
    print(f"Maze size: {maze.shape[1]}x{maze.shape[2]}")
    print(f"Solution path length: {len(solution_path)}")
    
    # Encode maze
    encoded = adapter.encode_maze(maze)
    print(f"\nEncoded shape: {encoded.shape}")
    
    # Forward pass through HRM
    with torch.no_grad():
        output = model(encoded)
    
    print(f"HRM output shape: {output['output'].shape}")
    
    # Decode path
    actions, predicted_paths = adapter.decode_path(
        output['output'],
        maze
    )
    
    print(f"\nAction logits shape: {actions.shape}")
    print(f"Number of predicted paths: {len(predicted_paths)}")
    
    if predicted_paths[0]:
        print(f"Predicted path length: {len(predicted_paths[0])}")
        print(f"Start position: {predicted_paths[0][0]}")
        print(f"End position: {predicted_paths[0][-1]}")
        
        # Check if reached goal
        goal_pos = tuple(torch.where(maze[0] == 3)[0].tolist())
        if goal_pos and len(goal_pos) == 2:
            goal_pos = (goal_pos[0][0], goal_pos[1][0]) if isinstance(goal_pos[0], list) else goal_pos
            reached_goal = predicted_paths[0][-1] == goal_pos
            print(f"Reached goal (untrained): {reached_goal}")
    
    return model, adapter

def demo_model_info():
    """Show model information"""
    print("="*60)
    print("MODEL ARCHITECTURE INFO")
    print("="*60)
    
    model = HierarchicalReasoningModel(
        input_dim=512,
        output_dim=512,
        high_hidden_dim=512,
        low_hidden_dim=256,
        num_iterations=3
    )
    
    param_counts = model.get_param_count()
    
    print("Parameter counts by component:")
    for key, value in param_counts.items():
        if key != 'total_millions':
            print(f"  {key}: {value:,}")
    
    print(f"\nTotal parameters: {param_counts['total_millions']:.2f}M")
    print(f"Target limit: 27M")
    print(f"Efficiency: {(param_counts['total_millions'] / 27 * 100):.1f}% of limit")
    
    # Show a sample forward pass
    print("\nSample forward pass:")
    x = torch.randn(2, 10, 512)  # batch=2, seq=10, dim=512
    
    with torch.no_grad():
        output = model(x, return_intermediates=True)
    
    print(f"  Input shape: {x.shape}")
    print(f"  Output shape: {output['output'].shape}")
    print(f"  High-level states: {output['high_states'].shape}")
    print(f"  Low-level states: {output['low_states'].shape}")
    print(f"  Number of iterations: {len(output['intermediates'])}")
    
    return model

def main():
    """Run all demos"""
    print("HIERARCHICAL REASONING MODEL - DEMO")
    print("="*60)
    print("This demo shows the HRM architecture in action")
    print("Note: Model is untrained, so outputs are random")
    print("="*60)
    
    # Show model info
    model = demo_model_info()
    
    # Demo Sudoku
    sudoku_model, sudoku_adapter = demo_sudoku()
    
    # Demo Maze
    maze_model, maze_adapter = demo_maze()
    
    print("\n" + "="*60)
    print("DEMO COMPLETE")
    print("="*60)
    print("\nTo train the model, run:")
    print("  python -m src.hrm.training.train --task sudoku --epochs 100")
    print("  python -m src.hrm.training.train --task maze --epochs 100")

if __name__ == "__main__":
    main()