progen2_sample.py

"""proto_tools/tools/causal_models/progen2/progen2_sample.py.

ProGen2 sampling tool.
"""

import logging
from typing import Any, Literal, cast

from pydantic import Field

from proto_tools.tools.causal_models.shared_data_models import (
    CausalModelSample,
    CausalModelSampleConfig,
    CausalModelSampleInput,
    CausalModelSampleOutput,
)
from proto_tools.tools.tool_registry import tool
from proto_tools.utils import (
    ConfigField,
    ToolInstance,
)
from proto_tools.utils.device import RemoteDevice

logger = logging.getLogger(__name__)

PROGEN2_MODEL_CHECKPOINTS = Literal[
    "progen2-small",
    "progen2-medium",
    "progen2-base",
    "progen2-oas",
    "progen2-large",
    "progen2-BFD90",
    "progen2-xlarge",
]


# ============================================================================
# Data Models
# ============================================================================
ProGen2SampleInput = CausalModelSampleInput


class ProGen2Sample(CausalModelSample):
    """One generated protein sequence and its per-position logits.

    Attributes:
        sequence (str): The generated protein sequence.
        logits (list[list[float]] | None): Per-position logits for this sequence
            (shape: [generated_len, vocab_size]).
    """

    logits: list[list[float]] | None = Field(
        default=None,
        title="Logits",
        description="Per-position logits for this generated sequence",
    )


class ProGen2SampleOutput(CausalModelSampleOutput):
    """Output from ProGen2 protein sequence generation.

    Attributes:
        results (list[ProGen2Sample]): One generated protein sequence per prompt, with its logits.
    """

    results: list[ProGen2Sample] = Field(  # type: ignore[assignment]
        title="Results",
        description="Generated protein sequences with their optional logits, one per prompt",
    )


# Config:
class ProGen2SampleConfig(CausalModelSampleConfig):
    """Configuration for ProGen2 protein sequence sampling.

    ProGen2 is an autoregressive protein language model with sizes from 151M to 6B
    parameters and specialized variants for antibody (OAS) and broader protein
    families (BFD90).

    Attributes:
        prepend_prompt (bool): Include the input prompt at the start of each generated
            sequence; when ``False``, only newly generated tokens are returned.
        batch_size (int): Number of prompts to process simultaneously on GPU.
        model_checkpoint (PROGEN2_MODEL_CHECKPOINTS): ProGen2 weights variant.
            Sizes range from 151M (small) to 6B (xlarge).
        local_path (str | None): Override the default download with a local weights directory.
        temperature (float): Softmax temperature; lower values are more deterministic.
        top_p (float): Nucleus sampling threshold over per-position token probabilities.
        top_k (int): Top-k truncation; ``0`` disables and uses top-p only.
        max_new_tokens (int): Maximum number of new tokens to generate per prompt (excludes prompt).
        truncate_at_stop (bool): Truncate generated sequences at the first stop token.
        strip_special_tokens (bool): Strip ProGen2 start/stop sentinel tokens (``1``/``2``)
            from output.
        return_logits (bool): Include per-position logits in the output.
    """

    temperature: float = ConfigField(
        title="Temperature",
        default=0.2,
        gt=0.0,
        description="Softmax temperature for sampling; lower is more deterministic",
    )
    top_p: float = ConfigField(
        title="Top P",
        default=0.95,
        gt=0.0,
        le=1.0,
        description="Nucleus sampling threshold over per-position token probabilities",
    )
    model_checkpoint: PROGEN2_MODEL_CHECKPOINTS = ConfigField(
        default="progen2-large",
        title="Model Checkpoint",
        description="ProGen2 weights variant",
        reload_on_change=True,
    )
    local_path: str | None = ConfigField(
        default=None,
        title="Local Model Path",
        description="Override the default download with a local weights directory",
        reload_on_change=True,
    )
    top_k: int = ConfigField(
        default=0,
        ge=0,
        title="Top-k",
        description="Top-k truncation; 0 disables and uses top-p only",
    )
    max_new_tokens: int = ConfigField(
        default=256,
        ge=1,
        title="Max New Tokens",
        description="Maximum newly-generated tokens per prompt (excludes the prompt)",
    )
    truncate_at_stop: bool = ConfigField(
        default=True,
        title="Truncate at Stop",
        description="Truncate generated sequences at the first stop token",
    )
    strip_special_tokens: bool = ConfigField(
        title="Strip Special Tokens",
        default=True,
        description="Strip ProGen2 start/stop sentinel tokens from output",
    )
    return_logits: bool = ConfigField(
        title="Return Logits",
        default=False,
        description="Include per-position logits in the output (large; disable to save memory)",
    )

    @classmethod
    def minimal(cls, **kwargs: Any) -> "ProGen2SampleConfig":
        """Sample short sequences from the model's own distribution.

        The shipped ``temperature`` of 0.2 is deliberately conservative, and from a
        short prompt it concentrates the distribution enough to emit long
        single-residue repeats: separate draws pick the same token, so duplicate
        prompts come back identical. Sampling at 1.0 keeps the generated sequences
        varied, which is what the stochastic-tool test infrastructure exercises.
        """
        kwargs.setdefault("temperature", 1.0)
        kwargs.setdefault("max_new_tokens", 32)
        return cast("ProGen2SampleConfig", super().minimal(**kwargs))

    def remote_unsupported_reason(self, device: RemoteDevice) -> str | None:
        """A local weights directory (``local_path``) isn't present on a hosted worker."""
        if self.local_path:
            return f"local_path points to a local weights directory not available on device='{device}'. Unset it, or run locally with device='cpu'."
        return None


# ============================================================================
# Tool Implementation
# ============================================================================
def example_input() -> Any:
    """Minimal valid input for testing and examples."""
    return ProGen2SampleInput(prompts=["MKTL"])


def _build_progen2_samples(sequences: list[str], logits: list[list[list[float]]] | None) -> list[ProGen2Sample]:
    """Pair each generated sequence with its own logits, if the worker returned any."""
    return [
        ProGen2Sample(sequence=sequence, logits=None if logits is None else logits[i])
        for i, sequence in enumerate(sequences)
    ]


@tool(
    key="progen2-sample",
    label="ProGen2 Sampling",
    category="causal_models",
    input_class=ProGen2SampleInput,
    config_class=ProGen2SampleConfig,
    output_class=ProGen2SampleOutput,
    description="Sample protein sequences using ProGen2 language model",
    uses_gpu=True,
    stochastic=True,
    example_input=example_input,
    iterable_input_fields=["prompts"],
    iterable_output_field="results",
    max_chunk_size=32,
)
def run_progen2_sample(
    inputs: ProGen2SampleInput,
    config: ProGen2SampleConfig,
    instance: Any = None,
) -> ProGen2SampleOutput:
    """Generate protein sequences using ProGen2 autoregressive language model.

    Uses the ProGen2 protein language model to autoregressively generate protein
    sequences from prompt sequences. Supports local GPU execution with various
    sampling strategies.

    Args:
        inputs (ProGen2SampleInput): Validated input containing one or more protein
            prompt sequences. Prompts are tokenized as given; include ProGen2's
            start token '1' explicitly to condition generation on a sequence start.
        config (ProGen2SampleConfig): Validated ProGen2 sampling configuration specifying
            model variant, generation parameters (temperature, top-k, top-p),
            sequence length, and output processing options.

        instance (Any): Optional ToolInstance for subprocess execution.

    Returns:
        ProGen2SampleOutput: Structured output containing:
            - ``sequences``: List of generated protein sequences
            - Metadata about generation parameters and execution mode

    Examples:
        >>> # Basic protein sequence generation with explicit start token
        >>> inputs = ProGen2SampleInput(prompts=["1MKTL"])
        >>> config = ProGen2SampleConfig(max_new_tokens=100, temperature=0.2, top_p=0.95)
        >>> result = run_progen2_sample(inputs, config)
        >>> print(f"Generated: {result.sequences[0]}")

        >>> # Generate from a raw amino-acid prompt (tokenized as given)
        >>> inputs = ProGen2SampleInput(prompts=["MVLS"])
        >>> result = run_progen2_sample(inputs, config)

        >>> # Batch generation
        >>> inputs = ProGen2SampleInput(prompts=["1MKTL", "1MVLS", "1GSSGSSG"])
        >>> result = run_progen2_sample(inputs, config)
        >>> print(f"Generated {len(result.sequences)} sequences")

        >>> # Using antibody-specific model
        >>> config = ProGen2SampleConfig(model_checkpoint="progen2-oas", temperature=0.3)
        >>> result = run_progen2_sample(inputs, config)

    Note:
        - ProGen2 uses '1' as start token and '2' as stop token
        - Prompts are tokenized as given; prepend '1' to condition on a sequence start
        - Local execution runs inside a standalone venv via ToolInstance

    See Also:
        - HuggingFace: https://huggingface.co/hugohrban/
        - ProGen2-finetuning GitHub: https://github.com/hugohrban/ProGen2-finetuning
        - Original ProGen2 GitHub: https://github.com/enijkamp/progen2
    """
    logger.debug(f"Using local venv for ProGen2 sampling: {config.model_checkpoint}")
    result = ToolInstance.dispatch(
        "progen2",
        {
            "operation": "sample",
            "prompts": inputs.prompts,
            "model_checkpoint": config.model_checkpoint,
            "local_path": config.local_path,
            "temperature": config.temperature,
            "top_p": config.top_p,
            "top_k": config.top_k,
            "max_new_tokens": config.max_new_tokens,
            "truncate_at_stop": config.truncate_at_stop,
            "strip_special_tokens": config.strip_special_tokens,
            "prepend_prompt": config.prepend_prompt,
            "batch_size": config.batch_size,
            "device": config.device,
            "verbose": config.verbose,
            "return_logits": config.return_logits,
            "seed": config.seed,
        },
        instance=instance,
        config=config,
    )

    return ProGen2SampleOutput(
        metadata={
            "model_checkpoint": config.model_checkpoint,
            "temperature": config.temperature,
            "max_new_tokens": config.max_new_tokens,
        },
        results=_build_progen2_samples(result["sequences"], result.get("logits")),
    )