documentation

Author: finitearth

promptolution/callbacks.py

@@ -16,6 +16,15 @@ def on_train_end(self, logs=None):
 
 
 class LoggerCallback(Callback):
+    """
+    Callback for logging optimization progress.
+
+    This callback logs information about each step, epoch, and the end of training.
+
+    Attributes:
+        logger: The logger object to use for logging.
+        step (int): The current step number.
+    """
     def __init__(self, logger):
         # TODO check up whats up with logging leves
         self.logger = logger
@@ -36,6 +45,15 @@ def on_train_end(self, logs=None):
 
 
 class CSVCallback(Callback):
+    """
+    Callback for saving optimization progress to a CSV file.
+
+    This callback saves prompts and scores at each step to a CSV file.
+
+    Attributes:
+        path (str): The path to the CSV file.
+        step (int): The current step number.
+    """
     def __init__(self, path):
         # if dir does not exist
         if not os.path.exists(os.path.dirname(path)):
@@ -62,6 +80,15 @@ def on_train_end(self, logs=None):
 
 
 class BestPromptCallback(Callback):
+    """
+    Callback for tracking the best prompt during optimization.
+
+    This callback keeps track of the prompt with the highest score.
+
+    Attributes:
+        best_prompt (str): The prompt with the highest score so far.
+        best_score (float): The highest score achieved so far.
+    """
     def __init__(self):
         self.best_prompt = ""
         self.best_score = -99999
@@ -76,6 +103,14 @@ def get_best_prompt(self):
 
 
 class ProgressBarCallback(Callback):
+    """
+    Callback for displaying a progress bar during optimization.
+
+    This callback uses tqdm to display a progress bar that updates at each step.
+
+    Attributes:
+        pbar (tqdm): The tqdm progress bar object.
+    """
     def __init__(self, total_steps):
         self.pbar = tqdm(total=total_steps)
 

promptolution/config.py

@@ -4,6 +4,27 @@
 
 @dataclass
 class Config:
+    """
+    Configuration class for the promptolution library.
+
+    This class handles loading and parsing of configuration settings,
+    either from a config file or from keyword arguments.
+
+    Attributes:
+        task_name (str): Name of the task.
+        ds_path (str): Path to the dataset.
+        n_steps (int): Number of optimization steps.
+        optimizer (str): Name of the optimizer to use.
+        meta_prompt_path (str): Path to the meta prompt file.
+        meta_llms (str): Name of the meta language model.
+        downstream_llm (str): Name of the downstream language model.
+        evaluation_llm (str): Name of the evaluation language model.
+        init_pop_size (int): Initial population size. Defaults to 10.
+        logging_dir (str): Directory for logging. Defaults to "logs/run.csv".
+        experiment_name (str): Name of the experiment. Defaults to "experiment".
+        include_task_desc (bool): Whether to include task description. Defaults to False.
+        random_seed (int): Random seed for reproducibility. Defaults to 42.
+    """
     task_name: str
     ds_path: str
     n_steps: int

promptolution/llms/__init__.py

@@ -4,6 +4,25 @@
 
 
 def get_llm(model_id: str, *args, **kwargs):
+    """
+    Factory function to create and return a language model instance based on the provided model_id.
+
+    This function supports three types of language models:
+    1. DummyLLM: A mock LLM for testing purposes.
+    2. LocalLLM: For running models locally (identified by 'local' in the model_id).
+    3. APILLM: For API-based models (default if not matching other types).
+
+    Args:
+        model_id (str): Identifier for the model to use. Special cases:
+                        - "dummy" for DummyLLM
+                        - "local-{model_name}" for LocalLLM
+                        - Any other string for APILLM
+        *args: Variable length argument list passed to the LLM constructor.
+        **kwargs: Arbitrary keyword arguments passed to the LLM constructor.
+
+    Returns:
+        An instance of DummyLLM, LocalLLM, or APILLM based on the model_id.
+    """
     if model_id == "dummy":
         return DummyLLM(*args, **kwargs)
     if "local" in model_id:

promptolution/llms/api_llm.py

@@ -4,6 +4,8 @@
 import openai
 from logging import INFO, Logger
 
+from typing import List
+
 from langchain_anthropic import ChatAnthropic
 from langchain_community.chat_models.deepinfra import ChatDeepInfraException
 from langchain_core.messages import HumanMessage
@@ -17,6 +19,20 @@
 
 
 async def invoke_model(prompt, model, semaphore):
+    """
+    Asynchronously invoke a language model with retry logic.
+
+    Args:
+        prompt (str): The input prompt for the model.
+        model: The language model to invoke.
+        semaphore (asyncio.Semaphore): Semaphore to limit concurrent calls.
+
+    Returns:
+        str: The model's response content.
+
+    Raises:
+        ChatDeepInfraException: If all retry attempts fail.
+    """
     async with semaphore:
         max_retries = 100
         delay = 3
@@ -33,7 +49,30 @@ async def invoke_model(prompt, model, semaphore):
 
 
 class APILLM:
+    """
+    A class to interface with various language models through their respective APIs.
+
+    This class supports Claude (Anthropic), GPT (OpenAI), and LLaMA (DeepInfra) models.
+    It handles API key management, model initialization, and provides methods for
+    both synchronous and asynchronous inference.
+
+    Attributes:
+        model: The initialized language model instance.
+
+    Methods:
+        get_response: Synchronously get responses for a list of prompts.
+        _get_response: Asynchronously get responses for a list of prompts.
+    """
     def __init__(self, model_id: str):
+        """
+        Initialize the APILLM with a specific model.
+
+        Args:
+            model_id (str): Identifier for the model to use.
+
+        Raises:
+            ValueError: If an unknown model identifier is provided.
+        """
         if "claude" in model_id:
             ANTHROPIC_API_KEY = open("anthropictoken.txt", "r").read()
             self.model = ChatAnthropic(model=model_id, api_key=ANTHROPIC_API_KEY)
@@ -46,7 +85,21 @@ def __init__(self, model_id: str):
         else:
             raise ValueError(f"Unknown model: {model_id}")
 
-    def get_response(self, prompts: list[str]) -> list[str]:
+    def get_response(self, prompts: List[str]) -> List[str]:
+        """
+        Synchronously get responses for a list of prompts.
+
+        This method includes retry logic for handling connection errors and rate limits.
+
+        Args:
+            prompts (list[str]): List of input prompts.
+
+        Returns:
+            list[str]: List of model responses.
+
+        Raises:
+            requests.exceptions.ConnectionError: If max retries are exceeded.
+        """
         max_retries = 100
         delay = 3
         attempts = 0
@@ -74,6 +127,18 @@ def get_response(self, prompts: list[str]) -> list[str]:
     async def _get_response(
         self, prompts: list[str], max_concurrent_calls=200
     ) -> list[str]:  # TODO change name of method
+        """
+        Asynchronously get responses for a list of prompts.
+
+        This method uses a semaphore to limit the number of concurrent API calls.
+
+        Args:
+            prompts (list[str]): List of input prompts.
+            max_concurrent_calls (int): Maximum number of concurrent API calls allowed.
+
+        Returns:
+            list[str]: List of model responses.
+        """
         semaphore = asyncio.Semaphore(max_concurrent_calls)  # Limit the number of concurrent calls
         tasks = []
 

promptolution/llms/base_llm.py

@@ -5,19 +5,60 @@
 
 
 class BaseLLM(ABC):
+    """
+    Abstract base class for Language Models in the promptolution library.
+
+    This class defines the interface that all concrete LLM implementations should follow.
+
+    Methods:
+        get_response: An abstract method that should be implemented by subclasses
+                      to generate responses for given prompts.
+    """
     def __init__(self, *args, **kwargs):
         pass
 
     @abstractmethod
     def get_response(self, prompts: List[str]) -> List[str]:
+        """
+        Generate responses for the given prompts.
+
+        This method should be implemented by subclasses to define how
+        the LLM generates responses.
+
+        Args:
+            prompts (List[str]): A list of input prompts.
+
+        Returns:
+            List[str]: A list of generated responses corresponding to the input prompts.
+        """
         pass
 
 
 class DummyLLM(BaseLLM):
+    """
+    A dummy implementation of the BaseLLM for testing purposes.
+
+    This class generates random responses for given prompts, simulating
+    the behavior of a language model without actually performing any
+    complex natural language processing.
+    """
     def __init__(self, *args, **kwargs):
         pass
 
     def get_response(self, prompts: str) -> str:
+        """
+        Generate random responses for the given prompts.
+
+        This method creates silly, random responses enclosed in <prompt> tags.
+        It's designed for testing and demonstration purposes.
+
+        Args:
+            prompts (str or List[str]): Input prompt(s). If a single string is provided,
+                                        it's converted to a list containing that string.
+
+        Returns:
+            List[str]: A list of randomly generated responses, one for each input prompt.
+        """
         if isinstance(prompts, str):
             prompts = [prompts]
         results = []

promptolution/llms/deepinfra.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-
+# TODO delete?
 from typing import (
     Any,
     AsyncIterator,

promptolution/llms/local_llm.py

@@ -9,7 +9,30 @@
 
 
 class LocalLLM:
+    """
+    A class for running language models locally using the Hugging Face Transformers library.
+
+    This class sets up a text generation pipeline with specified model parameters
+    and provides a method to generate responses for given prompts.
+
+    Attributes:
+        pipeline (transformers.Pipeline): The text generation pipeline.
+
+    Methods:
+        get_response: Generate responses for a list of prompts.
+    """
     def __init__(self, model_id: str, batch_size=8):
+        """
+        Initialize the LocalLLM with a specific model.
+
+        Args:
+            model_id (str): The identifier of the model to use (e.g., "gpt2", "facebook/opt-1.3b").
+            batch_size (int, optional): The batch size for text generation. Defaults to 8.
+
+        Note:
+            This method sets up a text generation pipeline with bfloat16 precision,
+            automatic device mapping, and specific generation parameters.
+        """
         self.pipeline = transformers.pipeline(
             "text-generation",
             model=model_id,
@@ -24,6 +47,19 @@ def __init__(self, model_id: str, batch_size=8):
         self.pipeline.tokenizer.padding_side = "left"
 
     def get_response(self, prompts: list[str]):
+        """
+        Generate responses for a list of prompts using the local language model.
+
+        Args:
+            prompts (list[str]): A list of input prompts.
+
+        Returns:
+            list[str]: A list of generated responses corresponding to the input prompts.
+
+        Note:
+            This method uses torch.no_grad() for inference to reduce memory usage.
+            It handles both single and batch inputs, ensuring consistent output format.
+        """
         with torch.no_grad():
             response = self.pipeline(prompts, pad_token_id=self.pipeline.tokenizer.eos_token_id)
 
@@ -32,3 +68,10 @@ def get_response(self, prompts: list[str]):
 
         response = [r["generated_text"] for r in response]
         return response
+    
+    def __del__(self):
+        try:
+            del self.pipeline
+            torch.cuda.empty_cache()
+        except Exception as e:
+            logger.warning(f"Error during LocalLLM cleanup: {e}")

promptolution/optimizers/__init__.py

@@ -4,6 +4,26 @@
 
 
 def get_optimizer(config, *args, **kwargs):
+    """
+    Factory function to create and return an optimizer instance based on the provided configuration.
+
+    This function selects and instantiates the appropriate optimizer class based on the
+    'optimizer' field in the config object. It supports three types of optimizers:
+    'dummy', 'evopromptde', and 'evopromptga'.
+
+    Args:
+        config: A configuration object that must have an 'optimizer' attribute.
+                For 'evopromptde', it should also have a 'donor_random' attribute.
+                For 'evopromptga', it should also have a 'selection_mode' attribute.
+        *args: Variable length argument list passed to the optimizer constructor.
+        **kwargs: Arbitrary keyword arguments passed to the optimizer constructor.
+
+    Returns:
+        An instance of the specified optimizer class.
+
+    Raises:
+        ValueError: If an unknown optimizer type is specified in the config.
+    """
     if config.optimizer == "dummy":
         return DummyOptimizer(*args, **kwargs)
     if config.optimizer == "evopromptde":