init

llms/__init__.py

@@ -0,0 +1,16 @@
+"""This module is adapt from https://github.com/zeno-ml/zeno-build"""
+from .providers.hf_utils import generate_from_huggingface_completion
+from .providers.openai_utils import (
+    generate_from_openai_chat_completion,
+    generate_from_openai_completion,
+)
+from .providers.anthropic_utils import generate_from_anthropic_completion
+from .utils import call_llm
+
+__all__ = [
+    "generate_from_openai_completion",
+    "generate_from_openai_chat_completion",
+    "generate_from_huggingface_completion",
+    "call_llm",
+    "generate_from_anthropic_completion"
+]

llms/lm_config.py

@@ -0,0 +1,61 @@
+"""Config for language models."""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class LMConfig:
+    """A config for a language model.
+
+    Attributes:
+        provider: The name of the API provider.
+        model: The name of the model.
+        model_cls: The Python class corresponding to the model, mostly for
+             Hugging Face transformers.
+        tokenizer_cls: The Python class corresponding to the tokenizer, mostly
+            for Hugging Face transformers.
+        mode: The mode of the API calls, e.g., "chat" or "generation".
+    """
+
+    provider: str
+    model: str
+    model_cls: type | None = None
+    tokenizer_cls: type | None = None
+    mode: str | None = None
+    gen_config: dict[str, Any] = dataclasses.field(default_factory=dict)
+    cuda: str = '0'
+
+
+def construct_llm_config(args: argparse.Namespace) -> LMConfig:
+    llm_config = LMConfig(
+        provider=args.provider, model=args.model, mode=args.mode, cuda=args.cuda
+    )
+    if args.provider == "openai":
+        llm_config.gen_config["temperature"] = args.temperature
+        llm_config.gen_config["top_p"] = args.top_p
+        llm_config.gen_config["context_length"] = args.context_length
+        llm_config.gen_config["max_tokens"] = args.max_tokens
+        llm_config.gen_config["stop_token"] = args.stop_token
+        llm_config.gen_config["max_obs_length"] = args.max_obs_length
+        llm_config.gen_config["max_retry"] = args.max_retry
+    elif args.provider == "huggingface":
+        llm_config.gen_config["temperature"] = args.temperature
+        llm_config.gen_config["top_p"] = args.top_p
+        llm_config.gen_config["max_new_tokens"] = args.max_tokens
+        llm_config.gen_config["stop_sequences"] = (
+            [args.stop_token] if args.stop_token else None
+        )
+        llm_config.gen_config["max_obs_length"] = args.max_obs_length
+        llm_config.gen_config["model_endpoint"] = args.model_endpoint
+        llm_config.gen_config["max_retry"] = args.max_retry
+    elif args.provider == "anthropic":
+        llm_config.gen_config["max_obs_length"] = None
+        llm_config.gen_config["max_retry"] = args.max_retry
+    else:
+        raise NotImplementedError(f"provider {args.provider} not implemented")
+    return llm_config

llms/providers/anthropic_utils.py

@@ -0,0 +1,45 @@
+# Use the native inference API to send a text message to Anthropic Claude.
+
+import boto3
+import json
+
+def generate_from_anthropic_completion(
+    prompt: str,
+    engine: str,
+) -> str:
+    # Create a Bedrock Runtime client in the AWS Region of your choice.
+    client = boto3.client("bedrock-runtime", region_name="us-east-1")
+
+    # Set the model ID, e.g., Claude 3 Haiku.
+    model_id = "anthropic.claude-3-sonnet-20240229-v1:0"
+
+    # Format the request payload using the model's native structure.
+    native_request = {
+        "anthropic_version": "bedrock-2023-05-31",
+        "max_tokens": 1024,
+        "temperature": 0.5,
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": prompt}],
+            }
+        ],
+    }
+
+    # Convert the native request to JSON.
+    request = json.dumps(native_request)
+
+    try:
+        # Invoke the model with the request.
+        response = client.invoke_model(modelId=engine, body=request)
+
+    except Exception as e: # (ClientError, Exception) as e:
+        raise KeyError(f"ERROR: Can't invoke '{engine}'. Reason: {e}")
+
+    # Decode the response body.
+    model_response = json.loads(response["body"].read())
+
+    # Extract and print the response text.
+    response_text = model_response["content"][0]["text"]
+
+    return response_text

llms/providers/hf_utils.py

@@ -0,0 +1,21 @@
+from text_generation import Client  # type: ignore
+
+
+def generate_from_huggingface_completion(
+    prompt: str,
+    model_endpoint: str,
+    temperature: float,
+    top_p: float,
+    max_new_tokens: int,
+    stop_sequences: list[str] | None = None,
+) -> str:
+    client = Client(model_endpoint, timeout=60)
+    generation: str = client.generate(
+        prompt=prompt,
+        temperature=temperature,
+        top_p=top_p,
+        max_new_tokens=max_new_tokens,
+        stop_sequences=stop_sequences,
+    ).generated_text
+
+    return generation

llms/providers/openai_utils.py

@@ -0,0 +1,283 @@
+"""Tools to generate from OpenAI prompts.
+Adopted from https://github.com/zeno-ml/zeno-build/"""
+
+import asyncio
+import logging
+import os
+import random
+import time
+from typing import Any
+
+import aiolimiter
+import openai
+from openai import OpenAI
+from tqdm.asyncio import tqdm_asyncio
+
+
+def retry_with_exponential_backoff(  # type: ignore
+    func,
+    initial_delay: float = 1,
+    exponential_base: float = 2,
+    jitter: bool = True,
+    max_retries: int = 3,
+    errors: tuple[Any] = (openai.RateLimitError,),
+):
+    """Retry a function with exponential backoff."""
+
+    def wrapper(*args, **kwargs):  # type: ignore
+        # Initialize variables
+        num_retries = 0
+        delay = initial_delay
+
+        # Loop until a successful response or max_retries is hit or an exception is raised
+        while True:
+            try:
+                return func(*args, **kwargs)
+            # Retry on specified errors
+            except errors as e:
+                # Increment retries
+                num_retries += 1
+
+                # Check if max retries has been reached
+                if num_retries > max_retries:
+                    raise Exception(
+                        f"Maximum number of retries ({max_retries}) exceeded."
+                    )
+
+                # Increment the delay
+                delay *= exponential_base * (1 + jitter * random.random())
+                print(f"Retrying in {delay} seconds.")
+                # Sleep for the delay
+                time.sleep(delay)
+
+            # Raise exceptions for any errors not specified
+            except Exception as e:
+                raise e
+
+    return wrapper
+
+
+async def _throttled_openai_completion_acreate(
+    engine: str,
+    prompt: str,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    limiter: aiolimiter.AsyncLimiter,
+) -> dict[str, Any]:
+    async with limiter:
+        for _ in range(3):
+            try:
+                return await openai.Completion.acreate(  # type: ignore
+                    engine=engine,
+                    prompt=prompt,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    top_p=top_p,
+                )
+            except openai.RateLimitError:
+                logging.warning(
+                    "OpenAI API rate limit exceeded. Sleeping for 10 seconds."
+                )
+                await asyncio.sleep(10)
+            except openai.APIError as e:
+                logging.warning(f"OpenAI API error: {e}")
+                break
+        return {"choices": [{"message": {"content": ""}}]}
+
+
+async def agenerate_from_openai_completion(
+    prompts: list[str],
+    engine: str,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    context_length: int,
+    requests_per_minute: int = 300,
+) -> list[str]:
+    """Generate from OpenAI Completion API.
+
+    Args:
+        prompts: list of prompts
+        temperature: Temperature to use.
+        max_tokens: Maximum number of tokens to generate.
+        top_p: Top p to use.
+        context_length: Length of context to use.
+        requests_per_minute: Number of requests per minute to allow.
+
+    Returns:
+        List of generated responses.
+    """
+    if "OPENAI_API_KEY" not in os.environ:
+        raise ValueError(
+            "OPENAI_API_KEY environment variable must be set when using OpenAI API."
+        )
+    openai.api_key = os.environ["OPENAI_API_KEY"]
+    openai.organization = os.environ.get("OPENAI_ORGANIZATION", "")
+    limiter = aiolimiter.AsyncLimiter(requests_per_minute)
+    async_responses = [
+        _throttled_openai_completion_acreate(
+            engine=engine,
+            prompt=prompt,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            limiter=limiter,
+        )
+        for prompt in prompts
+    ]
+    responses = await tqdm_asyncio.gather(*async_responses)
+    return [x["choices"][0]["text"] for x in responses]
+
+
+@retry_with_exponential_backoff
+def generate_from_openai_completion(
+    prompt: str,
+    engine: str,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    stop_token: str | None = None,
+) -> str:
+    if "OPENAI_API_KEY" not in os.environ:
+        raise ValueError(
+            "OPENAI_API_KEY environment variable must be set when using OpenAI API."
+        )
+    openai.api_key = os.environ["OPENAI_API_KEY"]
+    openai.organization = os.environ.get("OPENAI_ORGANIZATION", "")
+    response = OpenAI().completions.create(  # type: ignore
+        prompt=prompt,
+        engine=engine,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        top_p=top_p,
+        stop=[stop_token],
+    )
+    answer: str = response["choices"][0]["text"]
+    return answer
+
+
+async def _throttled_openai_chat_completion_acreate(
+    model: str,
+    messages: list[dict[str, str]],
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    limiter: aiolimiter.AsyncLimiter,
+) -> dict[str, Any]:
+    async with limiter:
+        for _ in range(3):
+            try:
+                return await openai.ChatCompletion.acreate(  # type: ignore
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    top_p=top_p,
+                )
+            except openai.RateLimitError:
+                logging.warning(
+                    "OpenAI API rate limit exceeded. Sleeping for 10 seconds."
+                )
+                await asyncio.sleep(10)
+            except asyncio.exceptions.TimeoutError:
+                logging.warning("OpenAI API timeout. Sleeping for 10 seconds.")
+                await asyncio.sleep(10)
+            except openai.APIError as e:
+                logging.warning(f"OpenAI API error: {e}")
+                break
+        return {"choices": [{"message": {"content": ""}}]}
+
+
+async def agenerate_from_openai_chat_completion(
+    messages_list: list[list[dict[str, str]]],
+    engine: str,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    context_length: int,
+    requests_per_minute: int = 300,
+) -> list[str]:
+    """Generate from OpenAI Chat Completion API.
+
+    Args:
+        messages_list: list of message list
+        temperature: Temperature to use.
+        max_tokens: Maximum number of tokens to generate.
+        top_p: Top p to use.
+        context_length: Length of context to use.
+        requests_per_minute: Number of requests per minute to allow.
+
+    Returns:
+        List of generated responses.
+    """
+    if "OPENAI_API_KEY" not in os.environ:
+        raise ValueError(
+            "OPENAI_API_KEY environment variable must be set when using OpenAI API."
+        )
+    openai.api_key = os.environ["OPENAI_API_KEY"]
+    openai.organization = os.environ.get("OPENAI_ORGANIZATION", "")
+    limiter = aiolimiter.AsyncLimiter(requests_per_minute)
+    async_responses = [
+        _throttled_openai_chat_completion_acreate(
+            model=engine,
+            messages=message,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            limiter=limiter,
+        )
+        for message in messages_list
+    ]
+    responses = await tqdm_asyncio.gather(*async_responses)
+    return [x["choices"][0]["message"]["content"] for x in responses]
+
+
+@retry_with_exponential_backoff
+def generate_from_openai_chat_completion(
+    messages: list[dict[str, str]],
+    model: str,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    context_length: int,
+    stop_token: str | None = None,
+) -> str:
+    if "OPENAI_API_KEY" not in os.environ:
+        raise ValueError(
+            "OPENAI_API_KEY environment variable must be set when using OpenAI API."
+        )
+    openai.api_key = os.environ["OPENAI_API_KEY"]
+    openai.organization = os.environ.get("OPENAI_ORGANIZATION", "")
+
+    response = OpenAI().chat.completions.create(  # type: ignore
+        model=model,
+        messages=messages,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        top_p=top_p,
+        stop=[stop_token] if stop_token else None,
+    )
+    answer: str = response.choices[0].message.content
+    return answer
+
+
+@retry_with_exponential_backoff
+# debug only
+def fake_generate_from_openai_chat_completion(
+    messages: list[dict[str, str]],
+    model: str,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    context_length: int,
+    stop_token: str | None = None,
+) -> str:
+    if "OPENAI_API_KEY" not in os.environ:
+        raise ValueError(
+            "OPENAI_API_KEY environment variable must be set when using OpenAI API."
+        )
+    openai.api_key = os.environ["OPENAI_API_KEY"]
+    openai.organization = os.environ.get("OPENAI_ORGANIZATION", "")
+    answer = "Let's think step-by-step. This page shows a list of links and buttons. There is a search box with the label 'Search query'. I will click on the search box to type the query. So the action I will perform is \"click [60]\"."
+    return answer

llms/tokenizers.py

@@ -0,0 +1,29 @@
+from typing import Any
+
+import tiktoken
+from transformers import LlamaTokenizer  # type: ignore
+
+
+class Tokenizer(object):
+    def __init__(self, provider: str, model_name: str) -> None:
+        if provider == "openai":
+            self.tokenizer = tiktoken.encoding_for_model(model_name)
+        elif provider == "huggingface":
+            self.tokenizer = LlamaTokenizer.from_pretrained(model_name)
+            # turn off adding special tokens automatically
+            self.tokenizer.add_special_tokens = False  # type: ignore[attr-defined]
+            self.tokenizer.add_bos_token = False  # type: ignore[attr-defined]
+            self.tokenizer.add_eos_token = False  # type: ignore[attr-defined]
+        elif provider == "anthropic":
+            pass
+        else:
+            raise NotImplementedError
+
+    def encode(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text)
+
+    def decode(self, ids: list[int]) -> str:
+        return self.tokenizer.decode(ids)
+
+    def __call__(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text)

llms/utils.py

@@ -0,0 +1,67 @@
+import argparse
+from typing import Any
+from transformers import AutoTokenizer, AutoModel
+
+from llms import (
+    generate_from_huggingface_completion,
+    generate_from_openai_chat_completion,
+    generate_from_openai_completion,
+    lm_config,
+)
+
+APIInput = str | list[Any] | dict[str, Any]
+
+model = None
+tokenizer = None
+
+def call_llm(
+    lm_config: lm_config.LMConfig,
+    prompt: APIInput,
+) -> str:
+    global model
+    global tokenizer
+    
+    response: str
+    
+    if lm_config.provider == "openai":
+        if lm_config.mode == "chat":
+            assert isinstance(prompt, list)
+            response = generate_from_openai_chat_completion(
+                messages=prompt,
+                model=lm_config.model,
+                temperature=lm_config.gen_config["temperature"],
+                top_p=lm_config.gen_config["top_p"],
+                context_length=lm_config.gen_config["context_length"],
+                max_tokens=lm_config.gen_config["max_tokens"],
+                stop_token=None,
+            )
+        elif lm_config.mode == "completion":
+            assert isinstance(prompt, str)
+            response = generate_from_openai_completion(
+                prompt=prompt,
+                engine=lm_config.model,
+                temperature=lm_config.gen_config["temperature"],
+                max_tokens=lm_config.gen_config["max_tokens"],
+                top_p=lm_config.gen_config["top_p"],
+                stop_token=lm_config.gen_config["stop_token"],
+            )
+        else:
+            raise ValueError(
+                f"OpenAI models do not support mode {lm_config.mode}"
+            )
+    elif lm_config.provider == "huggingface":
+        assert isinstance(prompt, str)
+        response = generate_from_huggingface_completion(
+            prompt=prompt,
+            model_endpoint=lm_config.gen_config["model_endpoint"],
+            temperature=lm_config.gen_config["temperature"],
+            top_p=lm_config.gen_config["top_p"],
+            stop_sequences=lm_config.gen_config["stop_sequences"],
+            max_new_tokens=lm_config.gen_config["max_new_tokens"],
+        )
+    else:
+        raise NotImplementedError(
+            f"Provider {lm_config.provider} not implemented"
+        )
+
+    return response