rc-foundry 0.1.1__py3-none-any.whl

foundry/utils/weights.py

@@ -0,0 +1,271 @@
+"""Utils for loading weights from checkpoints."""
+
+import re
+from dataclasses import dataclass, field
+from enum import StrEnum, auto
+from os import PathLike
+
+import torch
+from beartype.typing import Pattern
+from torch import nn
+
+from foundry.utils.ddp import RankedLogger
+
+ranked_logger = RankedLogger(__name__, rank_zero_only=True)
+
+
+class WeightLoadingError(Exception):
+    """Exception raised when there's an error loading weights."""
+
+    pass
+
+
+class WeightLoadingPolicy(StrEnum):
+    """Policy for handling weights when loading checkpoints."""
+
+    # Always keep default initialization, regardless of whether the parameter is in the checkpoint or shapes match
+    REINIT = auto()
+
+    # Always zero-initialize, regardless of whether the parameter is in the checkpoint or shapes match
+    ZERO_INIT = auto()
+
+    # Copy from checkpoint only when shapes match exactly, otherwise error
+    COPY = auto()
+
+    # Copy from checkpoint if tensors are the same rank, padding with zeros if shapes don't match exectly
+    COPY_AND_ZERO_PAD = auto()
+
+
+class _PatternPolicyMixin:
+    """Mixin for handling glob-to-regex pattern compilation and matching for parameter policies.
+
+    Patterns can use the following wildcards:
+        - * matches any sequence of characters
+        - ? matches any single character
+        - [abc] matches any character in the brackets
+        - . matches a literal dot
+
+    Examples:
+        - "model.*.weight" matches any weight parameter in the model
+        - "model.encoder?.weight" matches encoder1.weight, encoder2.weight, etc.
+        - "model.encoder[12].weight" matches encoder1.weight and encoder2.weight
+        - "model.encoder.*.bias" matches any bias parameter in encoder submodules
+    """
+
+    _compiled_patterns: dict[Pattern, any]
+
+    @staticmethod
+    def _glob_to_regex(pattern: str) -> str:
+        # Convert glob pattern to regex string
+        return (
+            pattern.replace(".", r"\.")
+            .replace("*", ".*")
+            .replace("?", ".")
+            .replace("[", "[")
+            .replace("]", "]")
+        )
+
+    def _compile_patterns(self, policy_dict: dict[str, any]) -> dict[Pattern, any]:
+        compiled = {}
+        for pattern, value in list(policy_dict.items()):
+            if any(c in pattern for c in ["*", "?", "[", "]"]):
+                regex = self._glob_to_regex(pattern)
+                compiled[re.compile(f"^{regex}$")] = value
+        return compiled
+
+    def _get_policy_by_pattern(
+        self, param_name: str, policy_dict: dict[str, any], default: any
+    ) -> any:
+        # Exact match first
+        if policy_dict and param_name in policy_dict:
+            return policy_dict[param_name]
+        # Pattern match
+        for pattern, value in self._compiled_patterns.items():
+            if pattern.match(param_name):
+                return value
+        return default
+
+
+@dataclass
+class WeightLoadingConfig(_PatternPolicyMixin):
+    """Configuration for handling weights when loading a checkpoint."""
+
+    default_policy: WeightLoadingPolicy | str = WeightLoadingPolicy.COPY
+    fallback_policy: WeightLoadingPolicy | str = WeightLoadingPolicy.REINIT
+    param_policies: dict[str, WeightLoadingPolicy | str] = field(default_factory=dict)
+    _compiled_patterns: dict[Pattern, WeightLoadingPolicy] = field(
+        default_factory=dict, repr=False
+    )
+
+    def __post_init__(self):
+        if isinstance(self.default_policy, str):
+            self.default_policy = WeightLoadingPolicy(self.default_policy)
+        if isinstance(self.fallback_policy, str):
+            self.fallback_policy = WeightLoadingPolicy(self.fallback_policy)
+        for key, value in self.param_policies.items():
+            if isinstance(value, str):
+                self.param_policies[key] = WeightLoadingPolicy(value)
+        self._compiled_patterns = self._compile_patterns(self.param_policies)
+
+    def get_policy(self, param_name: str) -> WeightLoadingPolicy:
+        policy = self._get_policy_by_pattern(
+            param_name, self.param_policies, self.default_policy
+        )
+        assert isinstance(policy, WeightLoadingPolicy)
+        return policy
+
+
+@dataclass
+class ParameterFreezingConfig(_PatternPolicyMixin):
+    """Configuration for freezing model parameters after loading weights.
+
+    Allows specifying which parameters to freeze (set requires_grad=False) by exact name or pattern.
+    Patterns use glob-style wildcards (*, ?).
+
+    Attributes:
+        param_policies: Dict mapping parameter names or patterns to True (freeze) or False (do not freeze).
+        freeze_by_default: Whether to freeze parameters not matched by any pattern (default: False).
+    """
+
+    param_policies: dict[str, bool] = field(default_factory=dict)
+    freeze_by_default: bool = False
+    _compiled_patterns: dict[Pattern, bool] = field(default_factory=dict, repr=False)
+
+    def __post_init__(self):
+        self._compiled_patterns = self._compile_patterns(self.param_policies)
+
+    def is_frozen(self, param_name: str) -> bool:
+        """Get whether a parameter is frozen according to the config."""
+        is_frozen = self._get_policy_by_pattern(
+            param_name, self.param_policies, self.freeze_by_default
+        )
+        assert isinstance(is_frozen, bool)
+        return is_frozen
+
+
+def freeze_parameters_with_config(
+    model: nn.Module, config: ParameterFreezingConfig, verbose: bool = True
+) -> None:
+    """Freeze (set requires_grad=False) or unfreeze parameters according to config.
+
+    Args:
+        model: The model whose parameters to freeze/unfreeze.
+        config: ParameterFreezingConfig specifying which parameters to freeze.
+        verbose: Whether to log which parameters have non-default policies applied.
+    """
+    for name, param in model.named_parameters():
+        is_frozen = config.is_frozen(name)
+        param.requires_grad = not is_frozen
+
+        if is_frozen != config.freeze_by_default and verbose:
+            ranked_logger.info(f"Non-default freezing applied to {name}: {is_frozen}")
+
+
+def load_weights_with_policies(
+    model: nn.Module,
+    ckpt: dict[str, torch.Tensor],
+    config: WeightLoadingConfig | None = None,
+) -> dict:
+    """Load checkpoint weights into model according to the specified configuration.
+
+    Allows for partial loading of weights and zero-initialization of mismatched and arbitrary parameters.
+
+    Args:
+        model: The model to load weights INTO. By default, all model weights are re-initialized; we overwrite
+            with the checkpoint weights where appropriate
+        ckpt: Dictionary mapping parameter names to tensors (loaded from checkpoint on disk)
+        config: Configuration for handling weight loading. If None, uses default config
+    Returns:
+        dict: The updated state_dict (not loaded into model yet)
+    """
+    if config is None:
+        # (Initialize default config if not provided)
+        config = WeightLoadingConfig()
+
+    current_state = model.state_dict()
+    updated_state = {}  # We will update this with the new weights
+
+    def _apply_policy(
+        name: str,
+        current_param: torch.Tensor,
+        checkpoint_param: torch.Tensor | None,
+        policy: WeightLoadingPolicy,
+    ) -> torch.Tensor:
+        """Apply a weight loading policy and return the resulting tensor.
+
+        Raises WeightLoadingError for any policy application failures.
+        """
+        if policy == WeightLoadingPolicy.REINIT:
+            # Keep original initialization
+            return current_param
+
+        elif policy == WeightLoadingPolicy.ZERO_INIT:
+            # Zero-initialize
+            return torch.zeros_like(current_param)
+
+        elif policy == WeightLoadingPolicy.COPY:
+            # Must have checkpoint param and shapes must match
+            if checkpoint_param is None:
+                raise WeightLoadingError(f"Parameter '{name}' not found in checkpoint")
+            if current_param.shape != checkpoint_param.shape:
+                raise WeightLoadingError(
+                    f"Shape mismatch for '{name}': model {current_param.shape} vs checkpoint {checkpoint_param.shape}"
+                )
+            return checkpoint_param
+
+        elif policy == WeightLoadingPolicy.COPY_AND_ZERO_PAD:
+            # Must have checkpoint param and same number of dimensions
+            if checkpoint_param is None:
+                raise WeightLoadingError(f"Parameter '{name}' not found in checkpoint")
+            if len(current_param.shape) != len(checkpoint_param.shape):
+                raise WeightLoadingError(
+                    f"Different dimensions for '{name}': model {len(current_param.shape)}D vs checkpoint {len(checkpoint_param.shape)}D"
+                )
+
+            # Copy where shapes match, zero-init the rest
+            new_param = torch.zeros_like(current_param)
+            slices = tuple(
+                slice(0, min(d_ckpt, d_current))
+                for d_ckpt, d_current in zip(
+                    checkpoint_param.shape, current_param.shape
+                )
+            )
+            new_param[slices] = checkpoint_param[slices]
+            return new_param
+
+    # ... loop through all named parameters in the model
+    for name, current_param in current_state.items():
+        # Get the policy for this parameter
+        policy = config.get_policy(name)
+
+        # Get the corresponding parameter from the checkpoint
+        checkpoint_param = ckpt.get(name, None)
+
+        try:
+            # Try to apply the primary policy
+            result = _apply_policy(name, current_param, checkpoint_param, policy)
+            updated_state[name] = result
+        except WeightLoadingError as e:
+            # Primary policy failed, try fallback
+            ranked_logger.warning(
+                f"Failed to apply policy: '{policy}' to '{name}': {str(e)}. Falling back to policy: '{config.fallback_policy}'."
+            )
+            result = _apply_policy(
+                name, current_param, checkpoint_param, config.fallback_policy
+            )
+            updated_state[name] = result
+
+    return updated_state
+
+
+@dataclass
+class CheckpointConfig:
+    """Configuration for loading checkpoints.
+
+    TODO: Implement reset_scheduler and reset_ema
+    """
+
+    path: PathLike
+    reset_optimizer: bool = False
+    weight_loading_config: WeightLoadingConfig | None = None
+    parameter_freezing_config: ParameterFreezingConfig | None = None

foundry/version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.1.1'
+__version_tuple__ = version_tuple = (0, 1, 1)
+
+__commit_id__ = commit_id = None

foundry_cli/__init__.py

@@ -0,0 +1,3 @@
+"""Foundry CLI utilities."""
+
+__version__ = "0.1.0"

foundry_cli/download_checkpoints.py

@@ -0,0 +1,281 @@
+"""CLI for foundry model checkpoint installation and management."""
+
+import hashlib
+from pathlib import Path
+from typing import Optional
+from urllib.request import urlopen
+
+import rootutils
+import typer
+from dotenv import find_dotenv, load_dotenv, set_key
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    DownloadColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeRemainingColumn,
+    TransferSpeedColumn,
+)
+
+from foundry.inference_engines.checkpoint_registry import (
+    REGISTERED_CHECKPOINTS,
+    get_default_checkpoint_dir,
+)
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+load_dotenv(override=True)
+
+app = typer.Typer(help="Foundry model checkpoint installation utilities")
+console = Console()
+
+def download_file(url: str, dest: Path, verify_hash: Optional[str] = None) -> None:
+    """Download a file with progress bar and optional hash verification.
+
+    Args:
+        url: URL to download from
+        dest: Destination file path
+        verify_hash: Optional SHA256 hash to verify against
+
+    Raises:
+        ValueError: If hash verification fails
+    """
+    dest.parent.mkdir(parents=True, exist_ok=True)
+
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        DownloadColumn(),
+        TransferSpeedColumn(),
+        TimeRemainingColumn(),
+    ) as progress:
+        # Get file size
+        with urlopen(url) as response:
+            file_size = int(response.headers.get("Content-Length", 0))
+
+            task = progress.add_task(
+                f"Downloading {dest.name}", total=file_size, start=True
+            )
+
+            # Download with progress
+            hasher = hashlib.sha256() if verify_hash else None
+            with open(dest, "wb") as f:
+                while True:
+                    chunk = response.read(8192)
+                    if not chunk:
+                        break
+                    f.write(chunk)
+                    if hasher:
+                        hasher.update(chunk)
+                    progress.update(task, advance=len(chunk))
+
+    # Verify hash if provided
+    if verify_hash:
+        computed_hash = hasher.hexdigest()
+        if computed_hash != verify_hash:
+            dest.unlink()  # Remove corrupted file
+            raise ValueError(
+                f"Hash mismatch! Expected {verify_hash}, got {computed_hash}"
+            )
+        console.print("[green]✓[/green] Hash verification passed")
+
+
+def install_model(
+    model_name: str, checkpoint_dir: Path, force: bool = False
+) -> None:
+    """Install a single model checkpoint.
+
+    Args:
+        model_name: Name of the model (rfd3, rf3, mpnn)
+        checkpoint_dir: Directory to save checkpoints
+        force: Overwrite existing checkpoint if it exists
+    """
+    if model_name not in REGISTERED_CHECKPOINTS:
+        console.print(f"[red]Error:[/red] Unknown model '{model_name}'")
+        console.print(f"Available models: {', '.join(REGISTERED_CHECKPOINTS.keys())}")
+        raise typer.Exit(1)
+
+    checkpoint_info = REGISTERED_CHECKPOINTS[model_name]
+    dest_path = checkpoint_dir / checkpoint_info["filename"]
+
+    # Check if already exists
+    if dest_path.exists() and not force:
+        console.print(
+            f"[yellow]⚠[/yellow] {model_name} checkpoint already exists at {dest_path}"
+        )
+        console.print("Use --force to overwrite")
+        return
+
+    console.print(
+        f"[cyan]Installing {model_name}:[/cyan] {checkpoint_info['description']}"
+    )
+
+    try:
+        download_file(
+            checkpoint_info["url"], dest_path, checkpoint_info.get("sha256")
+        )
+        console.print(
+            f"[green]✓[/green] Successfully installed {model_name} to {dest_path}"
+        )
+    except Exception as e:
+        console.print(f"[red]✗[/red] Failed to install {model_name}: {e}")
+        raise typer.Exit(1)
+
+
+@app.command()
+def install(
+    models: list[str] = typer.Argument(
+        ...,
+        help="Models to install: 'all', 'rfd3', 'rf3', 'mpnn', or combination",
+    ),
+    checkpoint_dir: Optional[Path] = typer.Option(
+        None,
+        "--checkpoint-dir",
+        "-d",
+        help="Directory to save checkpoints (default: $FOUNDRY_CHECKPOINTS_DIR or ~/.foundry/checkpoints)",
+    ),
+    force: bool = typer.Option(
+        False, "--force", "-f", help="Overwrite existing checkpoints"
+    ),
+):
+    """Install model checkpoints for foundry.
+
+    Examples:
+
+        foundry install all
+
+        foundry install rfd3 rf3
+
+        foundry install proteinmpnn --checkpoint-dir ./checkpoints
+    """
+    # Determine checkpoint directory
+    if checkpoint_dir is None:
+        checkpoint_dir = get_default_checkpoint_dir()
+
+    console.print(f"[bold]Checkpoint directory:[/bold] {checkpoint_dir}")
+    console.print()
+
+    # Expand 'all' to all available models
+    if "all" in models:
+        models_to_install = ['rfd3', 'proteinmpnn', 'ligandmpnn', 'rf3']
+    else:
+        models_to_install = models
+
+    # Install each model
+    for model_name in models_to_install:
+        install_model(model_name, checkpoint_dir, force)
+        console.print()
+
+    set_key(
+        dotenv_path=find_dotenv(),
+        key_to_set='FOUNDRY_CHECKPOINTS_DIR',
+        value_to_set=str(checkpoint_dir),
+        export = False,
+    )
+    console.print(
+        f"Set checkpoint installation directory to: {checkpoint_dir}"
+    )
+
+    console.print("[bold green]Installation complete![/bold green]")
+
+
+@app.command(name="list")
+def list_models():
+    """List available model checkpoints."""
+    console.print("[bold]Available models:[/bold]\n")
+    for name, info in REGISTERED_CHECKPOINTS.items():
+        console.print(f"  [cyan]{name:8}[/cyan] - {info['description']}")
+
+
+@app.command()
+def show(
+    checkpoint_dir: Optional[Path] = typer.Option(
+        None,
+        "--checkpoint-dir",
+        "-d",
+        help="Checkpoint directory to show",
+    ),
+):
+    """Show installed checkpoints."""
+    if checkpoint_dir is None:
+        checkpoint_dir = get_default_checkpoint_dir()
+
+    if not checkpoint_dir.exists():
+        console.print(
+            f"[yellow]No checkpoints directory found at {checkpoint_dir}[/yellow]"
+        )
+        raise typer.Exit(0)
+
+    checkpoint_files = list(checkpoint_dir.glob("*.ckpt"))
+    if not checkpoint_files:
+        console.print(
+            f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]"
+        )
+        raise typer.Exit(0)
+
+    console.print(f"[bold]Installed checkpoints in {checkpoint_dir}:[/bold]\n")
+    total_size = 0
+    for ckpt in sorted(checkpoint_files):
+        size = ckpt.stat().st_size / (1024**3)  # GB
+        total_size += size
+        console.print(f"  {ckpt.name:30} {size:8.2f} GB")
+
+    console.print(f"\n[bold]Total:[/bold] {total_size:.2f} GB")
+
+
+@app.command()
+def clean(
+    checkpoint_dir: Optional[Path] = typer.Option(
+        None,
+        "--checkpoint-dir",
+        "-d",
+        help="Checkpoint directory to clean",
+    ),
+    confirm: bool = typer.Option(
+        True, "--confirm/--no-confirm", help="Ask for confirmation before deleting"
+    ),
+):
+    """Remove all downloaded checkpoints."""
+    if checkpoint_dir is None:
+        checkpoint_dir = get_default_checkpoint_dir()
+
+    if not checkpoint_dir.exists():
+        console.print(f"[yellow]No checkpoints found at {checkpoint_dir}[/yellow]")
+        raise typer.Exit(0)
+
+    # List files to delete
+    checkpoint_files = list(checkpoint_dir.glob("*.ckpt"))
+    if not checkpoint_files:
+        console.print(
+            f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]"
+        )
+        raise typer.Exit(0)
+
+    console.print("[bold]Files to delete:[/bold]")
+    total_size = 0
+    for ckpt in checkpoint_files:
+        size = ckpt.stat().st_size / (1024**3)  # GB
+        total_size += size
+        console.print(f"  {ckpt.name} ({size:.2f} GB)")
+
+    console.print(f"\n[bold]Total:[/bold] {total_size:.2f} GB")
+
+    # Confirm deletion
+    if confirm:
+        should_delete = typer.confirm("\nDelete these files?")
+        if not should_delete:
+            console.print("[yellow]Cancelled[/yellow]")
+            raise typer.Exit(0)
+
+    # Delete files
+    for ckpt in checkpoint_files:
+        ckpt.unlink()
+        console.print(f"[red]✗[/red] Deleted {ckpt.name}")
+
+    console.print("[green]✓[/green] Cleanup complete")
+
+
+if __name__ == "__main__":
+    app()

mpnn/__init__.py

@@ -0,0 +1 @@
+"""MPNN - ProteinMPNN and LigandMPNN implementations."""