bead 0.1.0__py3-none-any.whl

bead/items/adapters/openai.py

@@ -0,0 +1,323 @@
+"""OpenAI API adapter for item construction.
+
+This module provides a ModelAdapter implementation for OpenAI's API,
+supporting GPT models for various NLP tasks including log probability
+computation, embeddings, and natural language inference via prompting.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+
+try:
+    import openai
+except ImportError as e:
+    raise ImportError(
+        "openai package is required for OpenAI adapter. "
+        "Install it with: pip install openai"
+    ) from e
+
+from bead.items.adapters.api_utils import rate_limit, retry_with_backoff
+from bead.items.adapters.base import ModelAdapter
+from bead.items.cache import ModelOutputCache
+
+
+class OpenAIAdapter(ModelAdapter):
+    """Adapter for OpenAI API models.
+
+    Provides access to OpenAI's GPT models for language model operations,
+    embeddings, and prompted natural language inference.
+
+    Parameters
+    ----------
+    model_name : str
+        OpenAI model identifier (default: "gpt-3.5-turbo").
+    api_key : str | None
+        OpenAI API key. If None, uses OPENAI_API_KEY environment variable.
+    cache : ModelOutputCache | None
+        Cache for model outputs. If None, creates in-memory cache.
+    model_version : str
+        Model version for cache tracking (default: "latest").
+    embedding_model : str
+        Model to use for embeddings (default: "text-embedding-ada-002").
+
+    Attributes
+    ----------
+    model_name : str
+        OpenAI model identifier (e.g., "gpt-3.5-turbo", "gpt-4").
+    client : openai.OpenAI
+        OpenAI API client.
+    embedding_model : str
+        Model to use for embeddings (default: "text-embedding-ada-002").
+
+    Raises
+    ------
+    ValueError
+        If no API key is provided and OPENAI_API_KEY is not set.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "gpt-3.5-turbo",
+        api_key: str | None = None,
+        cache: ModelOutputCache | None = None,
+        model_version: str = "latest",
+        embedding_model: str = "text-embedding-ada-002",
+    ) -> None:
+        if cache is None:
+            cache = ModelOutputCache(backend="memory")
+
+        super().__init__(
+            model_name=model_name, cache=cache, model_version=model_version
+        )
+
+        # Get API key from parameter or environment
+        if api_key is None:
+            api_key = os.environ.get("OPENAI_API_KEY")
+            if api_key is None:
+                raise ValueError(
+                    "OpenAI API key must be provided via api_key parameter "
+                    "or OPENAI_API_KEY environment variable"
+                )
+
+        self.client = openai.OpenAI(api_key=api_key)
+        self.embedding_model = embedding_model
+
+    @retry_with_backoff(
+        max_retries=3,
+        initial_delay=1.0,
+        backoff_factor=2.0,
+        exceptions=(openai.APIError, openai.APIConnectionError, openai.RateLimitError),
+    )
+    @rate_limit(calls_per_minute=60)
+    def compute_log_probability(self, text: str) -> float:
+        """Compute log probability of text using OpenAI completions API.
+
+        Uses the completions API with logprobs to get token-level log probabilities
+        and sums them to get the total log probability.
+
+        Parameters
+        ----------
+        text : str
+            Text to compute log probability for.
+
+        Returns
+        -------
+        float
+            Log probability of the text (sum of token log probabilities).
+        """
+        # Check cache
+        cached = self.cache.get(
+            model_name=self.model_name, operation="log_probability", text=text
+        )
+        if cached is not None:
+            return float(cached)
+
+        # Call API
+        response = self.client.completions.create(
+            model=self.model_name,
+            prompt=text,
+            max_tokens=0,
+            echo=True,
+            logprobs=1,
+        )
+
+        # Sum token log probabilities
+        logprobs = response.choices[0].logprobs
+        if logprobs is None or logprobs.token_logprobs is None:
+            raise ValueError("API response did not include logprobs")
+
+        # Filter out None values (first token may have None)
+        token_logprobs = [lp for lp in logprobs.token_logprobs if lp is not None]
+        total_log_prob = sum(token_logprobs)
+
+        # Cache result
+        self.cache.set(
+            model_name=self.model_name,
+            operation="log_probability",
+            result=total_log_prob,
+            model_version=self.model_version,
+            text=text,
+        )
+
+        return float(total_log_prob)
+
+    def compute_perplexity(self, text: str) -> float:
+        """Compute perplexity of text.
+
+        Perplexity is computed as exp(-log_prob / num_tokens).
+
+        Parameters
+        ----------
+        text : str
+            Text to compute perplexity for.
+
+        Returns
+        -------
+        float
+            Perplexity of the text (must be positive).
+        """
+        # Check cache
+        cached = self.cache.get(
+            model_name=self.model_name, operation="perplexity", text=text
+        )
+        if cached is not None:
+            return float(cached)
+
+        # Get log probability
+        log_prob = self.compute_log_probability(text)
+
+        # Estimate number of tokens (rough approximation: 1 token ~ 4 chars)
+        num_tokens = max(1, len(text) // 4)
+
+        # Compute perplexity: exp(-log_prob / num_tokens)
+        perplexity = np.exp(-log_prob / num_tokens)
+
+        # Cache result
+        self.cache.set(
+            model_name=self.model_name,
+            operation="perplexity",
+            result=float(perplexity),
+            model_version=self.model_version,
+            text=text,
+        )
+
+        return float(perplexity)
+
+    @retry_with_backoff(
+        max_retries=3,
+        initial_delay=1.0,
+        backoff_factor=2.0,
+        exceptions=(openai.APIError, openai.APIConnectionError, openai.RateLimitError),
+    )
+    @rate_limit(calls_per_minute=60)
+    def get_embedding(self, text: str) -> np.ndarray:
+        """Get embedding vector for text using OpenAI embeddings API.
+
+        Parameters
+        ----------
+        text : str
+            Text to embed.
+
+        Returns
+        -------
+        np.ndarray
+            Embedding vector for the text.
+        """
+        # Check cache
+        cached = self.cache.get(
+            model_name=self.embedding_model, operation="embedding", text=text
+        )
+        if cached is not None:
+            return np.array(cached)
+
+        # Call API
+        response = self.client.embeddings.create(model=self.embedding_model, input=text)
+
+        embedding = np.array(response.data[0].embedding)
+
+        # Cache result
+        self.cache.set(
+            model_name=self.embedding_model,
+            operation="embedding",
+            result=embedding.tolist(),
+            model_version=self.model_version,
+            text=text,
+        )
+
+        return embedding
+
+    @retry_with_backoff(
+        max_retries=3,
+        initial_delay=1.0,
+        backoff_factor=2.0,
+        exceptions=(openai.APIError, openai.APIConnectionError, openai.RateLimitError),
+    )
+    @rate_limit(calls_per_minute=60)
+    def compute_nli(self, premise: str, hypothesis: str) -> dict[str, float]:
+        """Compute natural language inference scores via prompting.
+
+        Uses chat completions API with a prompt to classify the relationship
+        between premise and hypothesis.
+
+        Parameters
+        ----------
+        premise : str
+            Premise text.
+        hypothesis : str
+            Hypothesis text.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary with keys "entailment", "neutral", "contradiction"
+            mapping to probability scores.
+        """
+        # Check cache
+        cached = self.cache.get(
+            model_name=self.model_name,
+            operation="nli",
+            premise=premise,
+            hypothesis=hypothesis,
+        )
+        if cached is not None:
+            return dict(cached)
+
+        # Construct prompt
+        prompt = (
+            "Given the following premise and hypothesis, "
+            "determine the relationship between them.\n\n"
+            f"Premise: {premise}\n"
+            f"Hypothesis: {hypothesis}\n\n"
+            "Choose one of the following:\n"
+            "- entailment: The hypothesis is definitely true given the premise\n"
+            "- neutral: The hypothesis might be true given the premise\n"
+            "- contradiction: The hypothesis is definitely false given the premise\n\n"
+            "Respond with only one word: entailment, neutral, or contradiction."
+        )
+
+        # Call API
+        response = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0.0,
+            max_tokens=10,
+        )
+
+        # Parse response
+        answer = response.choices[0].message.content
+        if answer is None:
+            raise ValueError("API response did not include content")
+
+        answer = answer.strip().lower()
+
+        # Map to scores
+        scores: dict[str, float] = {
+            "entailment": 0.0,
+            "neutral": 0.0,
+            "contradiction": 0.0,
+        }
+
+        if "entailment" in answer:
+            scores["entailment"] = 1.0
+        elif "neutral" in answer:
+            scores["neutral"] = 1.0
+        elif "contradiction" in answer:
+            scores["contradiction"] = 1.0
+        else:
+            # Default to neutral if unclear
+            scores["neutral"] = 1.0
+
+        # Cache result
+        self.cache.set(
+            model_name=self.model_name,
+            operation="nli",
+            result=scores,
+            model_version=self.model_version,
+            premise=premise,
+            hypothesis=hypothesis,
+        )
+
+        return scores

bead/items/adapters/registry.py

@@ -0,0 +1,202 @@
+"""Model adapter registry for centralized adapter management.
+
+This module provides a registry for managing all model adapters,
+both local (HuggingFace) and API-based (OpenAI, Anthropic, etc.).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Unpack
+
+from typing_extensions import TypedDict
+
+if TYPE_CHECKING:
+    from bead.items.cache import ModelOutputCache
+
+
+class AdapterKwargs(TypedDict, total=False):
+    """Keyword arguments for adapter initialization."""
+
+    api_key: str
+    device: str
+    model_version: str
+    embedding_model: str
+    normalize_embeddings: bool
+    cache: ModelOutputCache
+
+
+from bead.items.adapters.base import ModelAdapter  # noqa: E402
+
+
+class ModelAdapterRegistry:
+    """Registry for all model adapters (local and API-based).
+
+    Provides centralized management of adapter types and instances,
+    with automatic instance caching to avoid redundant initialization.
+
+    Attributes
+    ----------
+    adapters : dict[str, type[ModelAdapter]]
+        Registered adapter classes keyed by adapter type name.
+    instances : dict[str, ModelAdapter]
+        Cached adapter instances keyed by unique identifier.
+    """
+
+    def __init__(self) -> None:
+        self.adapters: dict[str, type[ModelAdapter]] = {}
+        self.instances: dict[str, ModelAdapter] = {}
+
+    def register(self, name: str, adapter_class: type[ModelAdapter]) -> None:
+        """Register an adapter class.
+
+        Parameters
+        ----------
+        name : str
+            Unique name for the adapter type (e.g., "openai", "huggingface_lm").
+        adapter_class : type[ModelAdapter]
+            Adapter class to register (must inherit from ModelAdapter).
+
+        Raises
+        ------
+        ValueError
+            If adapter class does not inherit from ModelAdapter.
+        """
+        if not issubclass(adapter_class, ModelAdapter):  # type: ignore[misc]
+            raise ValueError(
+                f"Adapter class {adapter_class.__name__} must inherit from ModelAdapter"
+            )
+        self.adapters[name] = adapter_class
+
+    def get_adapter(
+        self, adapter_type: str, model_name: str, **kwargs: Unpack[AdapterKwargs]
+    ) -> ModelAdapter:
+        """Get or create adapter instance (with caching).
+
+        Creates a new adapter instance if not cached, otherwise returns
+        the cached instance. Instances are cached by adapter type and model name.
+
+        Parameters
+        ----------
+        adapter_type
+            Type of adapter (must be registered).
+        model_name
+            Model identifier for the adapter.
+        **kwargs
+            Additional keyword arguments to pass to adapter constructor
+            (api_key, device, model_version, embedding_model, etc.).
+
+        Returns
+        -------
+        ModelAdapter
+            Adapter instance (cached or newly created).
+
+        Raises
+        ------
+        ValueError
+            If adapter type is not registered.
+
+        Examples
+        --------
+        >>> registry = ModelAdapterRegistry()
+        >>> registry.register("openai", OpenAIAdapter)
+        >>> adapter = registry.get_adapter("openai", "gpt-4", api_key="...")
+        """
+        if adapter_type not in self.adapters:
+            raise ValueError(
+                f"Unknown adapter type: {adapter_type}. "
+                f"Available types: {list(self.adapters.keys())}"
+            )
+
+        # create cache key from adapter type and model name
+        cache_key = f"{adapter_type}:{model_name}"
+
+        # return cached instance if available
+        if cache_key in self.instances:
+            return self.instances[cache_key]
+
+        # create new instance
+        adapter_class = self.adapters[adapter_type]
+        adapter = adapter_class(model_name=model_name, **kwargs)  # type: ignore[misc]
+
+        # cache and return
+        self.instances[cache_key] = adapter
+        return adapter
+
+    def clear_cache(self) -> None:
+        """Clear all cached adapter instances.
+
+        Useful for testing or when you want to force recreation of adapters
+        with different parameters.
+        """
+        self.instances.clear()
+
+    def list_adapters(self) -> list[str]:
+        """List all registered adapter types.
+
+        Returns
+        -------
+        list[str]
+            List of registered adapter type names.
+        """
+        return list(self.adapters.keys())
+
+
+# Create default registry with all built-in adapters
+default_registry = ModelAdapterRegistry()
+
+# Register HuggingFace adapters
+try:
+    from bead.items.adapters.huggingface import (
+        HuggingFaceLanguageModel,
+        HuggingFaceMaskedLanguageModel,
+        HuggingFaceNLI,
+    )
+
+    default_registry.register("huggingface_lm", HuggingFaceLanguageModel)
+    default_registry.register("huggingface_mlm", HuggingFaceMaskedLanguageModel)
+    default_registry.register("huggingface_nli", HuggingFaceNLI)
+except ImportError:
+    # HuggingFace adapters not available (missing dependencies)
+    pass
+
+# Register sentence transformers
+try:
+    from bead.items.adapters.sentence_transformers import HuggingFaceSentenceTransformer
+
+    default_registry.register("sentence_transformer", HuggingFaceSentenceTransformer)
+except ImportError:
+    # Sentence transformers not available
+    pass
+
+# Register API adapters (these are optional)
+try:
+    from bead.items.adapters.openai import OpenAIAdapter
+
+    default_registry.register("openai", OpenAIAdapter)
+except ImportError:
+    # OpenAI not available
+    pass
+
+try:
+    from bead.items.adapters.anthropic import AnthropicAdapter
+
+    default_registry.register("anthropic", AnthropicAdapter)
+except ImportError:
+    # Anthropic not available
+    pass
+
+try:
+    from bead.items.adapters.google import GoogleAdapter
+
+    default_registry.register("google", GoogleAdapter)
+except ImportError:
+    # Google not available
+    pass
+
+try:
+    from bead.items.adapters.togetherai import TogetherAIAdapter
+
+    default_registry.register("togetherai", TogetherAIAdapter)
+except ImportError:
+    # Together AI not available
+    pass