feat: v0.1.0 Refactor & Packaging

.dockerignore

@@ -0,0 +1,65 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+*.egg-info/
+dist/
+build/
+.eggs/
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+env/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Git
+.git/
+.gitignore
+.gitmodules
+.gitattributes
+
+# Documentation
+docs/_build/
+*.md
+!README.md
+
+# Data and checkpoints (too large for image)
+data/
+checkpoints/
+wandb/
+cache/
+
+# Tests
+tests/
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+
+# CI/CD
+.github/
+.pre-commit-config.yaml
+
+# Notebooks
+*.ipynb
+.ipynb_checkpoints/
+
+# Temporary files
+*.log
+*.tmp
+.DS_Store
+Thumbs.db
+
+# Package files
+*.tar.gz
+*.whl

.github/workflows/ci.yml

@@ -0,0 +1,44 @@
+name: CI
+
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+
+jobs:
+  lint-and-test:
+    name: Lint & Test (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v5
+
+      - name: Set up uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.13"
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: uv sync --extra dev
+
+      - name: Lint check Python:${{ matrix.python-version }}
+        run: uv run ruff check
+
+      - name: Tests Python:${{ matrix.python-version }}
+        run: uv run pytest
+
+      # TODO: Setup CodeCov
+      # - name: Upload coverage reports to Codecov
+      #   uses: codecov/codecov-action@v5
+      #   with:
+      #     token: ${{ secrets.CODECOV_TOKEN }}
+      #     slug: sapientinc/HRM

.gitignore

@@ -166,4 +166,9 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+#.idea/
+
+CLAUDE.md
+INSTRUCTIONS.md
+.ruff_cache/
+.claude/

.python-version

@@ -0,0 +1 @@
+3.12

Dockerfile

@@ -0,0 +1,114 @@
+# Multi-stage Dockerfile for Hierarchical Reasoning Model
+# Supports both FlashAttention 2 (Ampere GPUs) and FlashAttention 3 (Hopper GPUs)
+
+# Build argument for FlashAttention version
+ARG FLASH_ATTN_VERSION=2
+
+# ============================================================================
+# Stage 1: Base Image with CUDA and System Dependencies
+# ============================================================================
+FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu24.04 AS base
+
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive \
+    CUDA_HOME=/usr/local/cuda-12.6 \
+    PATH=/usr/local/cuda-12.6/bin:$PATH \
+    LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64:$LD_LIBRARY_PATH \
+    PYTHONUNBUFFERED=1 \
+    PYTHON_VERSION=3.12
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    --no-install-recommends \
+    python${PYTHON_VERSION} \
+    python${PYTHON_VERSION}-dev \
+    python3-pip \
+    git \
+    wget \
+    build-essential \
+    ninja-build \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create symlink for python
+RUN ln -sf /usr/bin/python${PYTHON_VERSION} /usr/bin/python3 && \
+    ln -sf /usr/bin/python3 /usr/bin/python
+
+# ============================================================================
+# Stage 2: Install uv and Python Dependencies
+# ============================================================================
+FROM base AS builder
+
+# Copy uv from official image
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
+
+# Set working directory
+WORKDIR /app
+
+# Copy dependency files
+COPY pyproject.toml ./
+COPY README.md ./
+COPY LICENSE ./
+
+# Create src directory structure (required for build)
+RUN mkdir -p src/hierarchical_reasoning_model && \
+    touch src/hierarchical_reasoning_model/__init__.py
+
+# Install core dependencies
+RUN uv sync --frozen --no-dev
+
+# Install PyTorch with CUDA 12.6 support
+RUN uv pip install torch torchvision torchaudio \
+    --index-url https://download.pytorch.org/whl/cu126
+
+# Install FlashAttention based on version
+ARG FLASH_ATTN_VERSION
+RUN if [ "$FLASH_ATTN_VERSION" = "3" ]; then \
+        echo "Installing FlashAttention 3 for Hopper GPUs" && \
+        git clone https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention && \
+        cd /tmp/flash-attention/hopper && \
+        uv pip install . && \
+        rm -rf /tmp/flash-attention; \
+    else \
+        echo "Installing FlashAttention 2 for Ampere GPUs" && \
+        uv pip install flash-attn --no-build-isolation; \
+    fi
+
+# Install training dependencies
+RUN uv sync --frozen --extra train --extra data
+
+# ============================================================================
+# Stage 3: Final Runtime Image
+# ============================================================================
+FROM base AS runtime
+
+# Copy uv from builder
+COPY --from=builder /uv /uvx /bin/
+
+# Copy virtual environment from builder
+COPY --from=builder /app/.venv /app/.venv
+
+# Set working directory
+WORKDIR /app
+
+# Copy application code
+COPY src/ /app/src/
+COPY config/ /app/config/
+COPY pyproject.toml README.md LICENSE ./
+
+# Install the package in editable mode
+RUN . /app/.venv/bin/activate && \
+    uv pip install -e .
+
+# Set environment variables for runtime
+ENV PATH=/app/.venv/bin:$PATH \
+    PYTHONPATH=/app/src:$PYTHONPATH
+
+# Create directories for data and checkpoints
+RUN mkdir -p /app/data /app/checkpoints /app/wandb
+
+# Set default command
+CMD ["python", "-c", "import hierarchical_reasoning_model; print('HRM Docker image ready!')"]
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
+    CMD python -c "import torch; print('CUDA available:', torch.cuda.is_available())"