feat: add reason_first_program/sampling.py

reason_first_program/sampling.py

@@ -0,0 +1,473 @@
+"""
+Stage 1: Program Space Sampling
+
+Generate diverse valid implementations of a stub using multiple strategies:
+  - Direct sampling from LLMs at various temperatures
+  - SFS-inspired scattering (2411.05010): diversify via textual gradient directions
+  - Multi-model heterogeneous sampling (AlgoDiv finding: diversity requires multiple models)
+  - Concept-guided sampling: steer toward specific concept regions
+
+Supports both API-based models (OpenAI, Anthropic, HF Inference) and local models.
+"""
+
+from __future__ import annotations
+
+import re
+import time
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from reason_first_program.stub import Stub
+from reason_first_program.program_space import Program, ProgramSpace, execute_program
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SamplingConfig:
+    """Configuration for program sampling."""
+
+    n_samples: int = 100
+    temperatures: list[float] = field(
+        default_factory=lambda: [0.2, 0.6, 0.8, 1.0, 1.2]
+    )
+    models: list[str] = field(
+        default_factory=lambda: ["deepseek-coder"]
+    )
+    prompt_styles: list[str] = field(
+        default_factory=lambda: ["direct", "diverse"]
+    )
+    max_tokens: int = 1024
+    timeout_per_execution: float = 5.0
+    deduplicate: bool = True
+    filter_valid: bool = True
+
+
+class ModelBackend(ABC):
+    """Abstract backend for code generation."""
+
+    @abstractmethod
+    def generate(
+        self,
+        prompt: str,
+        temperature: float = 0.8,
+        max_tokens: int = 1024,
+        n: int = 1,
+    ) -> list[str]:
+        """Generate n completions for the given prompt."""
+        ...
+
+    @property
+    @abstractmethod
+    def model_id(self) -> str:
+        ...
+
+
+class HFInferenceBackend(ModelBackend):
+    """HuggingFace Inference API backend."""
+
+    def __init__(self, model_name: str = "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct", token: Optional[str] = None):
+        self._model_name = model_name
+        self._token = token
+
+    @property
+    def model_id(self) -> str:
+        return self._model_name
+
+    def generate(
+        self,
+        prompt: str,
+        temperature: float = 0.8,
+        max_tokens: int = 1024,
+        n: int = 1,
+    ) -> list[str]:
+        try:
+            from huggingface_hub import InferenceClient
+        except ImportError:
+            raise ImportError("pip install huggingface_hub")
+
+        client = InferenceClient(model=self._model_name, token=self._token)
+        results = []
+        for _ in range(n):
+            try:
+                response = client.text_generation(
+                    prompt,
+                    max_new_tokens=max_tokens,
+                    temperature=max(temperature, 0.01),
+                    do_sample=True,
+                )
+                results.append(response)
+            except Exception as e:
+                logger.warning(f"Generation failed: {e}")
+                continue
+        return results
+
+
+class OpenAIBackend(ModelBackend):
+    """OpenAI API backend."""
+
+    def __init__(self, model_name: str = "gpt-4o", api_key: Optional[str] = None):
+        self._model_name = model_name
+        self._api_key = api_key
+
+    @property
+    def model_id(self) -> str:
+        return self._model_name
+
+    def generate(
+        self,
+        prompt: str,
+        temperature: float = 0.8,
+        max_tokens: int = 1024,
+        n: int = 1,
+    ) -> list[str]:
+        try:
+            import openai
+        except ImportError:
+            raise ImportError("pip install openai")
+
+        client = openai.OpenAI(api_key=self._api_key)
+        try:
+            response = client.chat.completions.create(
+                model=self._model_name,
+                messages=[
+                    {"role": "system", "content": "You are an expert Python programmer. Output only the function body, no explanation."},
+                    {"role": "user", "content": prompt},
+                ],
+                temperature=temperature,
+                max_tokens=max_tokens,
+                n=n,
+            )
+            return [choice.message.content for choice in response.choices]
+        except Exception as e:
+            logger.warning(f"OpenAI generation failed: {e}")
+            return []
+
+
+class LocalModelBackend(ModelBackend):
+    """Local model backend using transformers."""
+
+    def __init__(self, model_name: str = "deepseek-ai/deepseek-coder-1.3b-instruct", device: str = "auto"):
+        self._model_name = model_name
+        self._device = device
+        self._pipeline = None
+
+    @property
+    def model_id(self) -> str:
+        return self._model_name
+
+    def _load(self):
+        if self._pipeline is None:
+            try:
+                from transformers import pipeline
+            except ImportError:
+                raise ImportError("pip install transformers torch")
+            self._pipeline = pipeline(
+                "text-generation",
+                model=self._model_name,
+                device_map=self._device,
+                trust_remote_code=True,
+            )
+
+    def generate(
+        self,
+        prompt: str,
+        temperature: float = 0.8,
+        max_tokens: int = 1024,
+        n: int = 1,
+    ) -> list[str]:
+        self._load()
+        results = []
+        for _ in range(n):
+            try:
+                out = self._pipeline(
+                    prompt,
+                    max_new_tokens=max_tokens,
+                    temperature=max(temperature, 0.01),
+                    do_sample=True,
+                    return_full_text=False,
+                )
+                results.append(out[0]["generated_text"])
+            except Exception as e:
+                logger.warning(f"Local generation failed: {e}")
+        return results
+
+
+def _extract_function_body(raw_output: str, stub: Stub) -> Optional[str]:
+    """
+    Extract a clean function body from LLM output.
+    Handles markdown code blocks, extra commentary, etc.
+    """
+    text = raw_output.strip()
+
+    # Remove markdown code fences
+    code_block = re.search(r"```(?:python)?\s*\n(.*?)```", text, re.DOTALL)
+    if code_block:
+        text = code_block.group(1).strip()
+
+    # If the output contains a full function def, extract it
+    func_match = re.search(
+        rf"def\s+{re.escape(stub.name)}\s*\(.*?\).*?:\s*\n(.*)",
+        text,
+        re.DOTALL,
+    )
+    if func_match:
+        text = func_match.group(1)
+
+    # Remove any leading/trailing non-code lines
+    lines = text.split("\n")
+    code_lines = []
+    in_code = False
+    for line in lines:
+        stripped = line.strip()
+        if stripped and not stripped.startswith("#") and not in_code:
+            in_code = True
+        if in_code or stripped.startswith("#"):
+            code_lines.append(line)
+
+    if not code_lines:
+        return None
+
+    return "\n".join(code_lines)
+
+
+def _build_full_source(body: str, stub: Stub) -> str:
+    """Reconstruct full function source from body and stub signature."""
+    # Extract just the def line from the stub source
+    for line in stub.source.split("\n"):
+        if line.strip().startswith("def "):
+            def_line = line
+            break
+    else:
+        def_line = f"def {stub.name}{stub.signature}:"
+
+    # Ensure proper indentation of body
+    indented_body = "\n".join(
+        f"    {line}" if line.strip() else line for line in body.split("\n")
+    )
+    return f"{def_line}\n{indented_body}"
+
+
+class ProgramSampler:
+    """
+    Basic program sampler: generates completions from a single backend.
+    """
+
+    def __init__(self, backend: ModelBackend, config: Optional[SamplingConfig] = None):
+        self.backend = backend
+        self.config = config or SamplingConfig()
+
+    def sample(self, stub: Stub) -> ProgramSpace:
+        """Sample programs for a stub and return a ProgramSpace."""
+        space = ProgramSpace(stub)
+        samples_per_config = max(
+            1,
+            self.config.n_samples
+            // (len(self.config.temperatures) * len(self.config.prompt_styles)),
+        )
+
+        for temp in self.config.temperatures:
+            for style in self.config.prompt_styles:
+                prompt = stub.to_completion_prompt(style=style)
+                logger.info(
+                    f"Sampling {samples_per_config} programs "
+                    f"(temp={temp}, style={style}, model={self.backend.model_id})"
+                )
+
+                raw_outputs = self.backend.generate(
+                    prompt=prompt,
+                    temperature=temp,
+                    max_tokens=self.config.max_tokens,
+                    n=samples_per_config,
+                )
+
+                for raw in raw_outputs:
+                    body = _extract_function_body(raw, stub)
+                    if body is None:
+                        continue
+
+                    full_source = _build_full_source(body, stub)
+                    program = Program(
+                        source=body,
+                        full_source=full_source,
+                        stub_id=stub.stub_id,
+                        model_id=self.backend.model_id,
+                        metadata={
+                            "temperature": temp,
+                            "prompt_style": style,
+                        },
+                    )
+
+                    # Execute and validate
+                    if stub.test_inputs:
+                        program = execute_program(
+                            program, stub, stub.test_inputs,
+                            timeout_seconds=self.config.timeout_per_execution,
+                        )
+
+                    space.add(program)
+
+        # Post-processing
+        if self.config.deduplicate:
+            space = space.deduplicate_syntactic()
+        if self.config.filter_valid and stub.test_inputs:
+            space = space.filter_valid()
+
+        return space
+
+
+class DiverseSampler:
+    """
+    Diverse program sampler using multiple backends and SFS-inspired scattering.
+
+    Key insight from AlgoDiv (2503.00691): combining solutions from heterogeneous
+    models increases algorithmic diversity more than any single-model technique.
+    """
+
+    def __init__(
+        self,
+        backends: list[ModelBackend],
+        config: Optional[SamplingConfig] = None,
+    ):
+        self.backends = backends
+        self.config = config or SamplingConfig()
+
+    def sample(self, stub: Stub) -> ProgramSpace:
+        """Sample from all backends and merge into a single ProgramSpace."""
+        space = ProgramSpace(stub)
+        samples_per_backend = max(1, self.config.n_samples // len(self.backends))
+
+        for backend in self.backends:
+            backend_config = SamplingConfig(
+                n_samples=samples_per_backend,
+                temperatures=self.config.temperatures,
+                models=[backend.model_id],
+                prompt_styles=self.config.prompt_styles,
+                max_tokens=self.config.max_tokens,
+                timeout_per_execution=self.config.timeout_per_execution,
+                deduplicate=False,  # We'll deduplicate at the end
+                filter_valid=False,
+            )
+            sampler = ProgramSampler(backend, backend_config)
+            backend_space = sampler.sample(stub)
+
+            for program in backend_space.programs:
+                space.add(program)
+
+            logger.info(
+                f"Backend {backend.model_id}: generated {len(backend_space)} programs"
+            )
+
+        # Post-processing across all backends
+        if self.config.deduplicate:
+            space = space.deduplicate_syntactic()
+        if self.config.filter_valid and stub.test_inputs:
+            space = space.filter_valid()
+
+        logger.info(
+            f"DiverseSampler: {len(space)} total programs "
+            f"({len(space.valid_programs)} valid)"
+        )
+        return space
+
+    def sample_with_scattering(
+        self, stub: Stub, n_directions: int = 5
+    ) -> ProgramSpace:
+        """
+        SFS-inspired scattering (2411.05010): first discover diverse algorithmic
+        directions, then sample implementations along each direction.
+        """
+        # Phase 1: Discover algorithmic directions
+        scout_backend = self.backends[0]
+        direction_prompt = (
+            f"Consider this Python function stub:\n\n"
+            f"```python\n{stub.source}\n```\n\n"
+            f"{stub.constraints.to_prompt_context()}\n\n"
+            f"List {n_directions} fundamentally different algorithmic approaches "
+            f"to implement this function. For each, give a short name and 1-sentence "
+            f"description. Format: '1. NAME: description'"
+        )
+        direction_outputs = scout_backend.generate(
+            direction_prompt, temperature=0.7, n=1
+        )
+
+        directions = []
+        if direction_outputs:
+            for line in direction_outputs[0].split("\n"):
+                line = line.strip()
+                if line and line[0].isdigit():
+                    # Extract direction name
+                    match = re.match(r"\d+\.\s*(.+?)(?::|$)", line)
+                    if match:
+                        directions.append(match.group(1).strip())
+
+        if not directions:
+            directions = [
+                "iterative approach",
+                "recursive approach",
+                "functional/map-reduce approach",
+                "optimized in-place approach",
+                "library-heavy approach",
+            ]
+
+        logger.info(f"Discovered {len(directions)} algorithmic directions: {directions}")
+
+        # Phase 2: Sample along each direction
+        space = ProgramSpace(stub)
+        samples_per_direction = max(
+            1, self.config.n_samples // (len(directions) * len(self.backends))
+        )
+
+        for direction in directions:
+            directed_prompt = (
+                f"Complete this Python function using the following approach: "
+                f"**{direction}**\n\n"
+                f"```python\n{stub.source}\n```\n\n"
+                f"{stub.constraints.to_prompt_context()}\n\n"
+                f"Only output the function body. Use the {direction} approach."
+            )
+
+            for backend in self.backends:
+                for temp in self.config.temperatures:
+                    raw_outputs = backend.generate(
+                        directed_prompt,
+                        temperature=temp,
+                        max_tokens=self.config.max_tokens,
+                        n=samples_per_direction,
+                    )
+
+                    for raw in raw_outputs:
+                        body = _extract_function_body(raw, stub)
+                        if body is None:
+                            continue
+
+                        full_source = _build_full_source(body, stub)
+                        program = Program(
+                            source=body,
+                            full_source=full_source,
+                            stub_id=stub.stub_id,
+                            model_id=backend.model_id,
+                            metadata={
+                                "temperature": temp,
+                                "direction": direction,
+                                "prompt_style": "scattered",
+                            },
+                        )
+
+                        if stub.test_inputs:
+                            program = execute_program(
+                                program, stub, stub.test_inputs,
+                                timeout_seconds=self.config.timeout_per_execution,
+                            )
+
+                        space.add(program)
+
+        # Post-processing
+        if self.config.deduplicate:
+            space = space.deduplicate_syntactic()
+        if self.config.filter_valid and stub.test_inputs:
+            space = space.filter_valid()
+
+        return space