config.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

"""
Configuration module for xLean paths.

This module provides centralized path configuration to make the codebase
portable and open-source friendly. Paths can be configured via:
1. Environment variables (highest priority)
2. A config file at ~/.xlean/config.json
3. Default values (fallback)

## Quick Setup

Set environment variables:

```bash
export XLEAN_DATA_DIR="/path/to/your/data"
export XLEAN_CHECKPOINT_DIR="/path/to/your/checkpoints"
export XLEAN_TOKENIZER_DIR="/path/to/your/tokenizers"
export XLEAN_OUTPUT_DIR="/path/to/your/output"
export XLEAN_REPO_DIR="/path/to/lean/repos"  # for sketching
```

Or create `~/.xlean/config.json`:

```json
{
  "data_dir": "/path/to/your/data",
  "checkpoint_dir": "/path/to/your/checkpoints",
  "tokenizer_dir": "/path/to/your/tokenizers",
  "output_dir": "/path/to/your/output"
}
```

## Environment Variables

| Variable              | Description                          | Default              |
|-----------------------|--------------------------------------|----------------------|
| XLEAN_DATA_DIR        | Base directory for data files        | ~/xlean_data         |
| XLEAN_CHECKPOINT_DIR  | Base directory for model checkpoints | ~/xlean_checkpoints  |
| XLEAN_TOKENIZER_DIR   | Directory containing tokenizer files | ~/xlean_tokenizers   |
| XLEAN_OUTPUT_DIR      | Default output directory for results | ~/xlean_output       |
| XLEAN_ELAN_BIN        | Path to elan/lean binaries           | ~/.elan/bin          |

## Expected Data Structure

```
$XLEAN_DATA_DIR/
├── minif2f/
│   ├── val.jsonl
│   └── test.jsonl
├── minif2f_curriculum/
│   └── val.jsonl
├── putnam/
│   └── val.jsonl
├── proofnet/
│   └── val.jsonl
└── shuffled/
    └── lean4_v10/

$XLEAN_CHECKPOINT_DIR/
├── lean4_v10/
│   └── 1e-4_acc/
└── Meta-Llama-3.3-70B-Instruct/

$XLEAN_TOKENIZER_DIR/
└── cl_toplang_128k.tiktoken
```

## Usage

```python
from xlean.config import paths, get_data_path, get_checkpoint_path

# Access configured directories
print(paths.data_dir)
print(paths.checkpoint_dir)

# Build paths
data_file = get_data_path("minif2f", "val.jsonl")
checkpoint = get_checkpoint_path("lean4_v10", "checkpoint.pt")
```
"""

import json
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional


def _get_default_data_dir() -> str:
    """Get default data directory."""
    return os.environ.get("XLEAN_DATA_DIR", str(Path.home() / "xlean_data"))


def _get_default_checkpoint_dir() -> str:
    """Get default checkpoint directory."""
    return os.environ.get("XLEAN_CHECKPOINT_DIR", str(Path.home() / "xlean_checkpoints"))


def _get_default_tokenizer_dir() -> str:
    """Get default tokenizer directory."""
    return os.environ.get("XLEAN_TOKENIZER_DIR", str(Path.home() / "xlean_tokenizers"))


def _get_default_output_dir() -> str:
    """Get default output directory."""
    return os.environ.get("XLEAN_OUTPUT_DIR", str(Path.home() / "xlean_output"))


def _get_elan_bin_dir() -> Optional[str]:
    """Get elan binary directory for PATH."""
    env_val = os.environ.get("XLEAN_ELAN_BIN")
    if env_val:
        return env_val
    default_elan = Path.home() / ".elan" / "bin"
    if default_elan.exists():
        return str(default_elan)
    return None


@dataclass
class XLeanPaths:
    """
    Centralized path configuration for xLean.

    All paths are configurable and default to sensible locations.
    Override via environment variables or by modifying this instance.
    """
    data_dir: str = field(default_factory=_get_default_data_dir)
    checkpoint_dir: str = field(default_factory=_get_default_checkpoint_dir)
    tokenizer_dir: str = field(default_factory=_get_default_tokenizer_dir)
    output_dir: str = field(default_factory=_get_default_output_dir)
    elan_bin: Optional[str] = field(default_factory=_get_elan_bin_dir)

    def __post_init__(self):
        """Load config from file if it exists."""
        config_file = Path.home() / ".xlean" / "config.json"
        if config_file.exists():
            with open(config_file) as f:
                config = json.load(f)
            for key, value in config.items():
                if hasattr(self, key) and not os.environ.get(f"XLEAN_{key.upper()}"):
                    setattr(self, key, value)

    def get_data_path(self, *parts: str) -> Path:
        """Get a path relative to the data directory."""
        return Path(self.data_dir).joinpath(*parts)

    def get_checkpoint_path(self, *parts: str) -> Path:
        """Get a path relative to the checkpoint directory."""
        return Path(self.checkpoint_dir).joinpath(*parts)

    def get_tokenizer_path(self, filename: str) -> Path:
        """Get the path to a tokenizer file."""
        return Path(self.tokenizer_dir) / filename

    def get_output_path(self, *parts: str) -> Path:
        """Get a path relative to the output directory."""
        return Path(self.output_dir).joinpath(*parts)

    def setup_lean_path(self) -> None:
        """Add elan binaries to PATH if available."""
        if self.elan_bin and self.elan_bin not in os.environ.get("PATH", ""):
            os.environ["PATH"] = f"{self.elan_bin}:{os.environ.get('PATH', '')}"

    def ensure_dirs_exist(self) -> None:
        """Create all configured directories if they don't exist."""
        for dir_path in [self.data_dir, self.checkpoint_dir, self.tokenizer_dir, self.output_dir]:
            Path(dir_path).mkdir(parents=True, exist_ok=True)

    @classmethod
    def from_env(cls) -> "XLeanPaths":
        """Create paths configuration from environment variables."""
        return cls()

    def to_dict(self) -> dict:
        """Export configuration as a dictionary."""
        return {
            "data_dir": self.data_dir,
            "checkpoint_dir": self.checkpoint_dir,
            "tokenizer_dir": self.tokenizer_dir,
            "output_dir": self.output_dir,
            "elan_bin": self.elan_bin,
        }

    def save_config(self, path: Optional[Path] = None) -> None:
        """Save configuration to a JSON file."""
        if path is None:
            path = Path.home() / ".xlean" / "config.json"
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)


# Global paths instance - import this in other modules
paths = XLeanPaths()


# Convenience functions for common paths
def get_tokenizer_path(filename: str = "cl_toplang_128k.tiktoken") -> str:
    """Get the full path to a tokenizer file."""
    return str(paths.get_tokenizer_path(filename))


def get_data_path(*parts: str) -> str:
    """Get a data path."""
    return str(paths.get_data_path(*parts))


def get_checkpoint_path(*parts: str) -> str:
    """Get a checkpoint path."""
    return str(paths.get_checkpoint_path(*parts))


def get_output_path(*parts: str) -> str:
    """Get an output path."""
    return str(paths.get_output_path(*parts))