wassname/Unsupervised-Elicitation
1
0
Fork 0
mirror of https://github.com/wassname/Unsupervised-Elicitation.git synced 2026-08-20 12:10:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

initial commit

Browse Source
This commit is contained in:
jiaxinwen
2025-06-10 18:56:05 +00:00
commit 2b1bc36340
34 changed files with 4394 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+155
View File
@@ -0,0 +1,155 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
ruff.toml
tensorboard/
# C extensions
*.so
.idea
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.vite/
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
SECRETS
data/
.DS_Store
pyrightconfig.json
*.db
prompt_history/*
prompt_history
wandb/
.hydra/
# API keys
.env
*api_key*
# results
src/results
results*
logs
README.md
+83
View File
@@ -0,0 +1,83 @@
## Unsupervised Elicitation of Language Models
We introduce a new unsupervised algorithm for eliciting skills from pretrained language models. This algorithm is competitive with training on human labels on common misconceptions (TruthfulQA), math (GSM8k-verification), and helpfulness reward modeling (Alpaca). Without supervision, we train a helpful chat assistant from the Haiku 3.5 base model that outperforms a similarly trained human-supervised baseline.
<p align="center">
<img width="100%" src="figures/llama_performance.png">
</p>
<p align="center">
<img width="100%" src="figures/claude_performance.png">
</p>
## Setup
### Environment
1. create conda environment: `conda env create -f env.yaml`
2. install package `pip install -e .`
### API for Pretrained Base Models
You should have access to an API for pretrained base models, which can return top-K (e.g. 20) logprobs.
Since most public api servers (e.g. openrouter) only support post-trained chat models, you probably need to deploy pretrained base models yourself. For example, we use vllm to deploy llama models in our experiments.
In particular, we highly recommend activating the `prefix caching` feature to accelerate the experiments, because our algorithm will create many API queries with similar prefixes.
### Secrets
You should create a file called SECRETS at the root of the repository with the following contents:
```
LLAMA_API_BASE=<your_api_base_url>
NYU_ORG=None
ARG_ORG=None
API_KEY=None
```
## Run
### ICM
<p align="center">
<img width="100%" src="figures/algorithm.png">
</p>
The main script is located in `src/experiments/ICM.py`
An example command for labeling truthfulQA data:
```
cd src/experiments
python ICM.py --testbed truthfulQA --alpha 50
```
Arguments:
- `--seed`: random seed
- `--alpha`: the coefficient for mutual predictability in our scoring function
- `--testbed`: name of the testbed, e.g., alpaca, truthfulqa, gsm8k
- `--model`: name of the pretrained base model, e.g., meta-llama/Llama-3.1-70B
- `--batch_size`: size of a minibatch when running ICM on large datasets that cannot be fit in to the context all at once[^1].
[^1]: Since ICM relies on in-context learning, it might not be able to fix all datapoints in the context at once. In our experiments, we split the whole dataset into $N$ batches (e.g., each batch consists of 256 datapoints) based on the context limit and data length, and run ICM independently on each batch.
- `--num_seed`: number of randomly labeled datapoints in the beginning.
- `--K`: max iteration
- `--consistency_fix_K`: max iteration for consistencyfix
- `--decay`: decay rate for simulating annealing
- `--initial_T`: initial temprature for simulated annealing
- `--final_T`: final temperature for simulated annealing
- `--scheduler`: decay scheduler for simulated annealing
### Iterative Fine-tuning
Instead of using the initial pretrained model ($M_0$) to label all $N$ batches, we do iterative fine-tuning:
- fine-tune the pretrained model on the first $j$ batches to obtain $M_j$
- use $M_j$ to label the $j+1$-th batch.
We use [axolotl](https://github.com/axolotl-ai-cloud/axolotl) for fine-tuning.
core/llm_api/__init__.py
View File
core/llm_api/anthropic_llm.py
+210
View File
@@ -0,0 +1,210 @@
import asyncio
import json
import logging
import os
import re
import time
from datetime import datetime
from traceback import format_exc
from typing import Optional, Union
import attrs
from anthropic import AsyncAnthropic
from anthropic.types import ContentBlock as AnthropicContentBlock
from termcolor import cprint
from core.llm_api.base_llm import PRINT_COLORS, LLMResponse, ModelAPIProtocol
from core.llm_api.openai_llm import OAIChatPrompt
ANTHROPIC_MODELS = {
"claude-instant-1",
"claude-2.0",
"claude-v1.3",
"claude-2.1",
"claude-3-opus-20240229",
"claude-3-sonnet-20240229",
"claude-3-haiku-20240307",
"claude-3-5-sonnet-20240620",
"claude-3-5-sonnet-20241022",
"claude-3-5-haiku-20241022"
}
LOGGER = logging.getLogger(__name__)
ACCEPTED_ARGS = [
"model",
"messages",
"max_tokens",
"system",
"max_tokens",
"stop_sequences",
"stream",
"temperature",
"top_p",
"top_k",
]
def extract_system_prompt(messages: list) -> str:
sys_prompt = ""
for message in messages:
if message["role"] == "system":
if sys_prompt:
raise ValueError(
"Multiple system messages found in the prompt. Only one is allowed."
)
sys_prompt = message["content"]
return sys_prompt
def transform_messages(messages: list) -> list:
_messages = []
for message in messages:
role = message["role"]
if role == "system":
continue
content = message["content"]
_messages.append({"role": role, "content": [{"type": "text", "text": content}]})
return _messages
def count_tokens(prompt: str) -> int:
return len(prompt.split())
def price_per_token(model_id: str) -> tuple[float, float]:
"""
Returns the (input token, output token) price for the given model id.
"""
return 0, 0
@attrs.define()
class AnthropicChatModel(ModelAPIProtocol):
num_threads: int
print_prompt_and_response: bool = False
client: AsyncAnthropic = attrs.field(
init=False, default=attrs.Factory(AsyncAnthropic)
)
available_requests: asyncio.BoundedSemaphore = attrs.field(init=False)
def __attrs_post_init__(self):
self.available_requests = asyncio.BoundedSemaphore(int(self.num_threads))
@staticmethod
def _create_prompt_history_file(prompt):
filename = f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}_prompt.txt"
with open(os.path.join("prompt_history", filename), "w") as f:
json_str = json.dumps(prompt, indent=4)
json_str = json_str.replace("\\n", "\n")
f.write(json_str)
return filename
@staticmethod
def _add_response_to_prompt_file(prompt_file, response):
with open(os.path.join("prompt_history", prompt_file), "a") as f:
f.write("\n\n======RESPONSE======\n\n")
json_str = json.dumps(response.to_dict(), indent=4)
json_str = json_str.replace("\\n", "\n")
f.write(json_str)
async def __call__(
self,
model_ids: list[str],
prompt: Union[str, OAIChatPrompt],
print_prompt_and_response: bool,
max_attempts: int,
**kwargs,
) -> list[LLMResponse]:
start = time.time()
assert (
len(model_ids) == 1
), "Anthropic implementation only supports one model at a time."
model_id = model_ids[0]
max_tokens = kwargs.pop("max_tokens", 2048)
LOGGER.debug(f"Making {model_id} call")
response: Optional[AnthropicContentBlock] = None
duration = None
kwargs = {k: v for k, v in kwargs.items() if k in ACCEPTED_ARGS}
system_prompt = extract_system_prompt(prompt)
raw_prompt = prompt
if "tool" in prompt[0]:
kwargs["tools"] = prompt[0]["tool"]
kwargs["tool_choice"] = prompt[0]["tool_choice"]
prompt = transform_messages(prompt)
# prompt_file = self._create_prompt_history_file([system_prompt] + prompt)
for i in range(max_attempts):
try:
async with self.available_requests:
api_start = time.time()
response = await self.client.messages.create(
messages=prompt,
model=model_id,
system=system_prompt,
max_tokens=max_tokens,
**kwargs,
)
api_duration = time.time() - api_start
except Exception as e:
error_info = f"Exception Type: {type(e).__name__}, Error Details: {str(e)}, Traceback: {format_exc()}"
LOGGER.warn(
f"Encountered API error: {error_info}.\nRetrying now. (Attempt {i})"
)
await asyncio.sleep(1.5**i)
else:
break
if response is None:
raise RuntimeError(
f"Failed to get a response from the API after {max_attempts} attempts."
)
num_context_tokens, num_completion_tokens = (
response.usage.input_tokens,
response.usage.output_tokens,
)
context_token_cost, completion_token_cost = price_per_token(model_id)
cost = (
num_context_tokens * context_token_cost
+ num_completion_tokens * completion_token_cost
)
duration = time.time() - start
LOGGER.debug(f"Completed call to {model_id} in {duration}s")
if "tools" in kwargs:
completion = response.content[0].input
else:
completion = response.content[0].text
llm_response = LLMResponse(
model_id=model_id,
completion=completion,
stop_reason=response.stop_reason,
duration=duration,
api_duration=api_duration,
cost=cost,
)
if self.print_prompt_and_response or print_prompt_and_response:
cprint("Prompt:", "white")
cprint("System: " + system_prompt, PRINT_COLORS["system"])
for message in prompt:
role = message["role"]
content = message["content"]
tag = "Human: " if role == "user" else "Assistant: "
cprint(tag, PRINT_COLORS[role], end="")
cprint(content, PRINT_COLORS[role])
cprint(f"Response ({llm_response.model_id}):", "white")
cprint(
f"{llm_response.completion}", PRINT_COLORS["assistant"], attrs=["bold"]
)
print()
return [{"prompt": raw_prompt, "response": llm_response.to_dict()}]
core/llm_api/base_llm.py
+137
View File
@@ -0,0 +1,137 @@
import json
import logging
from enum import Enum, auto
from typing import Dict, List, Optional, Protocol
import attrs
import numpy as np
from anthropic import AI_PROMPT, HUMAN_PROMPT
from pydantic import BaseModel
PRINT_COLORS = {"user": "cyan", "system": "magenta", "assistant": "light_green"}
LOGGER = logging.getLogger(__name__)
class PromptConfig(BaseModel):
partials: Dict[str, str] = {}
word_limit: Optional[int] = 100
messages: List[Dict[str, str]] = []
messages1: List[Dict[str, str]] = []
messages2: List[Dict[str, str]] = []
vars: Dict[str, str] = {}
class LanguageModelConfig(BaseModel):
model: str
temperature: float = 0.2
top_p: float = 1.0
max_tokens: Optional[int] = None
max_words: int = 10000
min_words: int = 0
num_candidates_per_completion: int = 1
timeout: int = 120
logit_bias: Optional[dict] = None
class StopReason(Enum):
MAX_TOKENS = auto()
STOP_SEQUENCE = auto()
TOOL_USE = auto()
@classmethod
def factory(cls, stop_reason: str) -> "StopReason":
"""
Parses the openai and anthropic stop reasons into a StopReason enum.
"""
if stop_reason in ["max_tokens", "length"]:
return cls.MAX_TOKENS
elif stop_reason in ["stop_sequence", "stop", "end_turn", "eos"]:
return cls.STOP_SEQUENCE
elif stop_reason in ['tool_use', "tool_calls"]:
return cls.TOOL_USE
raise ValueError(f"Invalid stop reason: {stop_reason}")
def __repr__(self):
return self.name
@attrs.frozen()
class LLMResponse:
model_id: str
completion: str
stop_reason: StopReason = attrs.field(converter=StopReason.factory)
cost: float
duration: Optional[float] = None
api_duration: Optional[float] = None
logprobs: Optional[list[dict[str, float]]] = None
def to_dict(self):
return {
"model_id": self.model_id,
"completion": self.completion,
"stop_reason": self.stop_reason.__repr__(), # Convert to some JSON-serializable format.
"duration": self.duration,
"api_duration": self.api_duration,
"cost": self.cost,
"logprobs": self.logprobs,
}
class ModelAPIProtocol(Protocol):
async def __call__(
self,
model_ids: list[str],
prompt,
print_prompt_and_response: bool,
max_attempts: int,
**kwargs,
) -> list[LLMResponse]:
raise NotImplementedError
def messages_to_single_prompt(messages) -> str:
if (
len(messages) >= 2
and messages[0]["role"] == "system"
and messages[1]["role"] == "user"
):
combined_content = messages[0]["content"] + " " + messages[1]["content"]
messages = [{"role": "user", "content": combined_content}] + messages[2:]
prompt = ""
for message in messages:
role = message["role"]
content = message["content"]
tag = AI_PROMPT if role == "assistant" else HUMAN_PROMPT
prompt += f"{tag} {content}"
if tag != AI_PROMPT:
prompt += f"{AI_PROMPT}"
return prompt.strip()
def convert_to_prob(log_prob: dict, tokens: list) -> tuple[float, float, float]:
logit1 = log_prob.get(tokens[0], None)
logit2 = log_prob.get(tokens[1], None)
if logit1 is None:
rating = -100
LOGGER.warning(
f"Missing token0 {tokens[0]} in log_prob, setting rating to -100.0"
)
else:
rating = logit1
if logit1 is None:
logit1 = -100
if logit2 is None:
logit2 = -100
return rating, logit1, logit2
def add_assistant_message(messages: list[dict], assistant_message: str):
last_role = messages[-1]["role"]
if last_role == "assistant":
messages[-1]["content"] += assistant_message
else:
messages.append({"role": "assistant", "content": assistant_message})
return messages
core/llm_api/llm.py
+316
View File
@@ -0,0 +1,316 @@
import asyncio
import json
import logging
import os
from collections import defaultdict
from itertools import chain
from pathlib import Path
from typing import Callable, Literal, Optional, Union
import attrs
from core.llm_api.anthropic_llm import ANTHROPIC_MODELS, AnthropicChatModel
from core.llm_api.base_llm import LLMResponse, ModelAPIProtocol
from core.llm_api.openai_llm import (
BASE_MODELS,
GPT_CHAT_MODELS,
OAIBasePrompt,
OAIChatPrompt,
OpenAIBaseModel,
OpenAIChatModel,
)
from core.utils import load_secrets
LOGGER = logging.getLogger(__name__)
@attrs.define()
class ModelAPI:
anthropic_num_threads: int = 2 # current redwood limit is 5
openai_fraction_rate_limit: float = attrs.field(
default=0.99, validator=attrs.validators.lt(1)
)
organization: str = "NYU_ORG"
print_prompt_and_response: bool = False
_openai_base: OpenAIBaseModel = attrs.field(init=False)
_openai_base_arg: OpenAIBaseModel = attrs.field(init=False)
_openai_chat: OpenAIChatModel = attrs.field(init=False)
_anthropic_chat: AnthropicChatModel = attrs.field(init=False)
running_cost: float = attrs.field(init=False, default=0)
model_timings: dict[str, list[float]] = attrs.field(init=False, default={})
model_wait_times: dict[str, list[float]] = attrs.field(init=False, default={})
def __attrs_post_init__(self):
secrets = load_secrets()
if self.organization is None:
self.organization = "NYU_ORG"
self._openai_base = OpenAIBaseModel(
frac_rate_limit=self.openai_fraction_rate_limit,
organization=secrets[self.organization],
print_prompt_and_response=self.print_prompt_and_response,
)
self._openai_base_arg = OpenAIBaseModel(
frac_rate_limit=self.openai_fraction_rate_limit,
organization=secrets["ARG_ORG"],
print_prompt_and_response=self.print_prompt_and_response,
)
self._openai_chat = OpenAIChatModel(
frac_rate_limit=self.openai_fraction_rate_limit,
organization=secrets[self.organization],
print_prompt_and_response=self.print_prompt_and_response,
)
self._anthropic_chat = AnthropicChatModel(
num_threads=self.anthropic_num_threads,
print_prompt_and_response=self.print_prompt_and_response,
)
Path("./prompt_history").mkdir(exist_ok=True)
@staticmethod
def _load_from_cache(save_file):
if not os.path.exists(save_file):
return None
else:
with open(save_file) as f:
cache = json.load(f)
return cache
async def call_single(
self,
model_ids: Union[str, list[str]],
prompt: Union[list[dict[str, str]], str],
max_tokens: int,
print_prompt_and_response: bool = False,
n: int = 1,
max_attempts_per_api_call: int = 10,
num_candidates_per_completion: int = 1,
# is_valid: Callable[[str], bool] = lambda _: True,
parse_fn=lambda _: True,
insufficient_valids_behaviour: Literal[
"error", "continue", "pad_invalids"
] = "error",
**kwargs,
) -> str:
assert n == 1, f"Expected a single response. {n} responses were requested."
responses = await self(
model_ids,
prompt,
max_tokens,
print_prompt_and_response,
n,
max_attempts_per_api_call,
num_candidates_per_completion,
parse_fn,
insufficient_valids_behaviour,
**kwargs,
)
assert len(responses) == 1, "Expected a single response."
return responses[0].completion
async def __call__(
self,
model_ids: Union[str, list[str]],
prompt: Union[list[dict[str, str]], str],
print_prompt_and_response: bool = False,
n: int = 1,
max_attempts_per_api_call: int = 50,
num_candidates_per_completion: int = 1,
parse_fn=None,
use_cache: bool = True,
file_sem: asyncio.Semaphore = None,
insufficient_valids_behaviour: Literal[
"error", "continue", "pad_invalids"
] = "error",
**kwargs,
) -> list[LLMResponse]:
"""
Make maximally efficient API requests for the specified model(s) and prompt.
Args:
model_ids: The model(s) to call. If multiple models are specified, the output will be sampled from the
cheapest model that has capacity. All models must be from the same class (e.g. OpenAI Base,
OpenAI Chat, or Anthropic Chat). Anthropic chat will error if multiple models are passed in.
Passing in multiple models could speed up the response time if one of the models is overloaded.
prompt: The prompt to send to the model(s). Type should match what's expected by the model(s).
max_tokens: The maximum number of tokens to request from the API (argument added to
standardize the Anthropic and OpenAI APIs, which have different names for this).
print_prompt_and_response: Whether to print the prompt and response to stdout.
n: The number of completions to request.
max_attempts_per_api_call: Passed to the underlying API call. If the API call fails (e.g. because the
API is overloaded), it will be retried this many times. If still fails, an exception will be raised.
num_candidates_per_completion: How many candidate completions to generate for each desired completion. n*num_candidates_per_completion completions will be generated, then is_valid is applied as a filter, then the remaining completions are returned up to a maximum of n.
parse_fn: post-processing on the generated response
save_path: cache path
use_cache: whether to load from the cache or overwrite it
"""
assert (
"max_tokens_to_sample" not in kwargs
), "max_tokens_to_sample should be passed in as max_tokens."
if isinstance(model_ids, str):
model_ids = [model_ids]
# # trick to double rate limit for most recent model only
def model_id_to_class(model_id: str) -> ModelAPIProtocol:
if model_id in ["gpt-4-base", "gpt-3.5-turbo-instruct"]:
return (
self._openai_base_arg
) # NYU ARG is only org with access to this model
elif model_id in BASE_MODELS:
return self._openai_base
elif model_id in GPT_CHAT_MODELS or "ft:gpt-3.5-turbo" in model_id:
return self._openai_chat
elif model_id in ANTHROPIC_MODELS:
return self._anthropic_chat
raise ValueError(f"Invalid model id: {model_id}")
model_classes = [model_id_to_class(model_id) for model_id in model_ids]
# assert model_classes == self._openai_base
# if model_classes == self._openai_base:
# assert "gpt" not in model_ids[0]
# kwargs['api_base'] = "https://5jfmglryfots6s-8000.proxy.runpod.net/v1"
if len(set(str(type(x)) for x in model_classes)) != 1:
raise ValueError("All model ids must be of the same type.")
max_tokens = (
kwargs.get("max_tokens") if kwargs.get("max_tokens") is not None else 2000
)
model_class = model_classes[0]
if isinstance(model_class, AnthropicChatModel):
kwargs["max_tokens_to_sample"] = max_tokens
else:
kwargs["max_tokens"] = max_tokens
# Check if current prompt has already been saved in the save file
# If so, directly return previous result
responses = None
if use_cache and kwargs.get("save_path") is not None:
try:
responses = self._load_from_cache(kwargs.get("save_path"))
except:
logging.error(f"invalid cache data: {kwargs.get('save_path')}")
# After loading cache, we do not directly return previous results,
# but continue running it through parse_fn and re-save it.
# This is because we may frequently update the parse_fn during development
if responses is None:
num_candidates = num_candidates_per_completion * n
if isinstance(model_class, AnthropicChatModel):
responses = list(
chain.from_iterable(
await asyncio.gather(
*[
model_class(
model_ids,
prompt,
print_prompt_and_response,
max_attempts_per_api_call,
**kwargs,
)
for _ in range(num_candidates)
]
)
)
)
else:
responses = await model_class(
model_ids,
prompt,
print_prompt_and_response,
max_attempts_per_api_call,
n=num_candidates,
**kwargs,
)
modified_responses = []
for response in responses:
self.running_cost += response["response"]["cost"]
if kwargs.get("metadata") is not None:
response["metadata"] = kwargs.get("metadata")
if parse_fn is not None:
response = parse_fn(response)
self.model_timings.setdefault(response["response"]["model_id"], []).append(
response["response"]["api_duration"]
)
self.model_wait_times.setdefault(
response["response"]["model_id"], []
).append(
response["response"]["duration"] - response["response"]["api_duration"]
)
modified_responses.append(response)
if kwargs.get("save_path") is not None:
if file_sem is not None:
async with file_sem:
with open(kwargs.get("save_path"), "w") as f:
json.dump(modified_responses, f, indent=2)
else:
with open(kwargs.get("save_path"), "w") as f:
json.dump(modified_responses, f, indent=2)
return modified_responses[:n]
def reset_cost(self):
self.running_cost = 0
async def demo():
model_api = ModelAPI(anthropic_num_threads=2, openai_fraction_rate_limit=0.99)
anthropic_requests = [
model_api(
"claude-3-5-sonnet-20240620",
[
{"role": "system", "content": "You are Claude."},
{"role": "user", "content": "who are you!"},
],
max_tokens=20,
print_prompt_and_response=False,
)
]
oai_chat_messages = [
[
{"role": "system", "content": "You are gpt-3.5-turbo."},
{"role": "user", "content": "who are you!"},
],
[
{
"role": "system",
"content": "You are gpt-4",
},
{"role": "user", "content": "who are you!"},
],
]
oai_chat_models = ["gpt-3.5-turbo-16k"]
oai_chat_requests = [
model_api(
oai_chat_models,
prompt=message,
max_tokens=16_000,
n=1,
print_prompt_and_response=False,
)
for message in oai_chat_messages
]
answer = await asyncio.gather(*anthropic_requests, *oai_chat_requests)
for responses in answer:
for i in responses:
print(i.completion)
print("=" * 100)
costs = defaultdict(int)
for responses in answer:
for response in responses:
costs[response.model_id] += response.cost
print("-" * 80)
print("Costs:")
for model_id, cost in costs.items():
print(f"{model_id}: ${cost}")
return answer
if __name__ == "__main__":
asyncio.run(demo())
core/llm_api/openai_llm.py
+631
View File
@@ -0,0 +1,631 @@
# %%
import asyncio
import json
import logging
import os
import random
import time
from datetime import datetime
from itertools import cycle
from traceback import format_exc
from typing import Optional, Union
import attrs
import openai
import requests
import tiktoken
from openai.openai_object import OpenAIObject as OpenAICompletion
from tenacity import retry, stop_after_attempt, wait_fixed
from termcolor import cprint
from core.llm_api.base_llm import (
PRINT_COLORS,
LLMResponse,
ModelAPIProtocol,
messages_to_single_prompt,
)
OAIChatPrompt = list[dict[str, str]]
OAIBasePrompt = Union[str, list[str]]
LOGGER = logging.getLogger(__name__)
def count_tokens(text: str) -> int:
return len(tiktoken.get_encoding("cl100k_base").encode(text))
def price_per_token(model_id: str) -> tuple[float, float]:
"""
Returns the (input token, output token) price for the given model id.
"""
if model_id == "gpt-4-1106-preview":
prices = 0.01, 0.03
elif model_id == "gpt-3.5-turbo-1106":
prices = 0.001, 0.002
elif model_id.startswith("gpt-4"):
prices = 0.03, 0.06
elif model_id.startswith("gpt-4-32k"):
prices = 0.06, 0.12
elif model_id.startswith("gpt-3.5-turbo-16k"):
prices = 0.003, 0.004
elif model_id.startswith("gpt-3.5-turbo"):
prices = 0.0015, 0.002
elif model_id == "davinci-002":
prices = 0.002, 0.002
elif model_id == "babbage-002":
prices = 0.0004, 0.0004
elif model_id == "text-davinci-003" or model_id == "text-davinci-002":
prices = 0.02, 0.02
elif "ft:gpt-3.5-turbo" in model_id:
prices = 0.012, 0.016
elif "llama" in model_id.lower() or "mixtral" in model_id.lower():
prices = 0.0015, 0.002
elif "o1" in model_id.lower():
prices = 0.01, 0.03
else:
prices = 0, 0
# raise ValueError(f"Invalid model id: {model_id}")
return tuple(price / 1000 for price in prices)
@attrs.define()
class Resource:
"""
A resource that is consumed over time and replenished at a constant rate.
"""
refresh_rate: float = (
attrs.field()
) # How many units of the resource are replenished per minute
value: float = attrs.field(init=False)
total: float = 0
throughput: float = 0
last_update_time: float = attrs.field(init=False, factory=time.time)
start_time: float = attrs.field(init=False, factory=time.time)
def __attrs_post_init__(self):
self.value = self.refresh_rate
def _replenish(self):
"""
Updates the value of the resource based on the time since the last update.
"""
curr_time = time.time()
self.value = min(
self.refresh_rate,
self.value + (curr_time - self.last_update_time) * self.refresh_rate / 60,
)
self.last_update_time = curr_time
self.throughput = self.total / (curr_time - self.start_time) * 60
def geq(self, amount: float) -> bool:
self._replenish()
return self.value >= amount
def consume(self, amount: float):
"""
Consumes the given amount of the resource.
"""
assert self.geq(
amount
), f"Resource does not have enough capacity to consume {amount} units"
self.value -= amount
self.total += amount
@attrs.define
class OpenAIModel(ModelAPIProtocol):
frac_rate_limit: float
organization: str
print_prompt_and_response: bool = False
model_ids: set[str] = attrs.field(init=False, default=attrs.Factory(set))
# rate limit
token_capacity: dict[str, Resource] = attrs.field(
init=False, default=attrs.Factory(dict)
)
request_capacity: dict[str, Resource] = attrs.field(
init=False, default=attrs.Factory(dict)
)
lock_add: asyncio.Lock = attrs.field(
init=False, default=attrs.Factory(asyncio.Lock)
)
lock_consume: asyncio.Lock = attrs.field(
init=False, default=attrs.Factory(asyncio.Lock)
)
@staticmethod
def _assert_valid_id(model_id: str):
raise NotImplementedError
@staticmethod
async def _get_dummy_response_header(model_id: str):
raise NotImplementedError
@staticmethod
def _count_prompt_token_capacity(prompt, **kwargs) -> int:
raise NotImplementedError
async def _make_api_call(self, prompt, model_id, **params) -> list[LLMResponse]:
raise NotImplementedError
@staticmethod
def _print_prompt_and_response(prompt, responses):
raise NotImplementedError
@staticmethod
def _create_prompt_history_file(prompt):
filename = f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}_prompt.txt"
with open(os.path.join("prompt_history", filename), "w") as f:
json_str = json.dumps(prompt, indent=4)
json_str = json_str.replace("\\n", "\n")
f.write(json_str)
return filename
@staticmethod
def _add_response_to_prompt_file(prompt_file, responses):
with open(os.path.join("prompt_history", prompt_file), "a") as f:
f.write("\n\n======RESPONSE======\n\n")
json_str = json.dumps(
[response.to_dict() for response in responses], indent=4
)
json_str = json_str.replace("\\n", "\n")
f.write(json_str)
async def add_model_id(self, model_id: str):
self._assert_valid_id(model_id)
if model_id in self.model_ids:
return
# make dummy request to get token and request capacity
model_metadata = await self._get_dummy_response_header(model_id)
token_capacity = int(model_metadata["x-ratelimit-limit-tokens"])
request_capacity = int(model_metadata["x-ratelimit-limit-requests"])
print(
f"got capacities for model {model_id}: {token_capacity}, {request_capacity}"
)
tokens_consumed = token_capacity - int(
model_metadata["x-ratelimit-remaining-tokens"]
)
requests_consumed = request_capacity - int(
model_metadata["x-ratelimit-remaining-requests"]
)
print(
f"consumed capacities for model {model_id}: {tokens_consumed}, {requests_consumed}"
)
token_cap = token_capacity * self.frac_rate_limit
request_cap = request_capacity * self.frac_rate_limit
if model_id in BASE_MODELS:
token_cap *= (
10000 # openai does not track token limit so we can increase it
)
print(f"setting cap for model {model_id}: {token_cap}, {request_cap}")
self.model_ids.add(model_id)
token_capacity = Resource(token_cap)
request_capacity = Resource(request_cap)
token_capacity.consume(min(token_cap, tokens_consumed))
request_capacity.consume(min(request_cap, requests_consumed))
self.token_capacity[model_id] = token_capacity
self.request_capacity[model_id] = request_capacity
async def __llama_call__(
self,
model_ids: list[str],
prompt,
print_prompt_and_response: bool,
max_attempts: int,
**kwargs,
) -> list[LLMResponse]:
kwargs = {
key: value
for key, value in kwargs.items()
if key not in ("save_path", "metadata")
}
start = time.time()
async def attempt_api_call():
api_base_list = [os.environ['LLAMA_API_BASE']]
kwargs["api_base"] = random.choice(api_base_list)
for model_id in cycle(model_ids):
return await asyncio.wait_for(
self._make_api_call(prompt, model_id, start, **kwargs),
timeout=100, # cloudflare has a 100-second limit for a connection to remain open: https://docs.runpod.io/pods/configuration/expose-ports
)
model_ids.sort(
key=lambda model_id: price_per_token(model_id)[0]
) # Default to cheapest model
model_id = model_ids[0]
prompt = self._process_prompt(prompt)
# prompt_file = self._create_prompt_history_file(prompt)
responses: Optional[list[LLMResponse]] = None
for i in range(max_attempts):
try:
responses = await attempt_api_call()
except Exception as e:
error_info = f"Exception Type: {type(e).__name__}, Error Details: {str(e)}, Traceback: {format_exc()}"
LOGGER.warn(
f"Encountered API error: {error_info}.\nRetrying now. (Attempt {i})"
)
await asyncio.sleep(1.5**i)
else:
break
if responses is None:
raise RuntimeError(
f"Failed to get a response from the API after {max_attempts} attempts."
)
if self.print_prompt_and_response or print_prompt_and_response:
self._print_prompt_and_response(prompt, responses)
end = time.time()
LOGGER.debug(f"Completed call to {model_id} in {end - start}s.")
return [
{"prompt": prompt, "response": response.to_dict()} for response in responses
]
async def __call__(
self,
model_ids: list[str],
prompt,
print_prompt_and_response: bool,
max_attempts: int,
**kwargs,
) -> list[LLMResponse]:
if "gpt" not in model_ids[0]:
return await self.__llama_call__(
model_ids, prompt, print_prompt_and_response, max_attempts, **kwargs
)
kwargs = {
key: value
for key, value in kwargs.items()
if key not in ("save_path", "metadata")
}
start = time.time()
async def attempt_api_call():
for model_id in cycle(model_ids):
async with self.lock_consume:
request_capacity, token_capacity = (
self.request_capacity[model_id],
self.token_capacity[model_id],
)
if request_capacity.geq(1) and token_capacity.geq(token_count):
request_capacity.consume(1)
token_capacity.consume(token_count)
else:
await asyncio.sleep(0.01)
continue # Skip this iteration if the condition isn't met
# Make the API call outside the lock
return await asyncio.wait_for(
self._make_api_call(prompt, model_id, start, **kwargs), timeout=120
)
model_ids.sort(
key=lambda model_id: price_per_token(model_id)[0]
) # Default to cheapest model
async with self.lock_add:
for model_id in model_ids:
await self.add_model_id(model_id)
if "tool" in prompt[0]:
kwargs["tools"] = prompt[0]["tool"]
if "response_format" in prompt[0]:
kwargs['response_format'] = prompt[0]['response_format']
prompt = self._process_prompt(prompt)
token_count = self._count_prompt_token_capacity(prompt, **kwargs)
assert (
max(self.token_capacity[model_id].refresh_rate for model_id in model_ids)
>= token_count
), "Prompt is too long for any model to handle."
# prompt_file = self._create_prompt_history_file(prompt)
responses: Optional[list[LLMResponse]] = None
for i in range(max_attempts):
try:
responses = await attempt_api_call()
except Exception as e:
error_info = f"Exception Type: {type(e).__name__}, Error Details: {str(e)}, Traceback: {format_exc()}"
LOGGER.warn(
f"Encountered API error: {error_info}.\nRetrying now. (Attempt {i})"
)
await asyncio.sleep(1.5**i)
else:
break
if responses is None:
raise RuntimeError(
f"Failed to get a response from the API after {max_attempts} attempts."
)
if self.print_prompt_and_response or print_prompt_and_response:
self._print_prompt_and_response(prompt, responses)
end = time.time()
LOGGER.debug(f"Completed call to {model_id} in {end - start}s.")
return [
{"prompt": prompt, "response": response.to_dict()} for response in responses
]
_GPT_4_MODELS = [
"gpt-4o",
"gpt-4",
"gpt-4-0314",
"gpt-4-0613",
"gpt-4-0125-preview",
"gpt-4-32k",
"gpt-4-32k-0314",
"gpt-4-32k-0613",
"gpt-4-1106-preview",
"gpt-4-turbo",
"gpt-4-turbo-preview",
"gpt-4-turbo-2024-04-09",
"gpt-4o-mini",
"gpt-4o-mini-2024-07-18",
"gpt-4o-2024-11-20",
"o1-preview-2024-09-12",
"o1-mini-2024-09-12",
"deepseek/deepseek-chat",
"meta-llama/llama-3.2-3b-instruct",
"meta-llama/llama-3.2-1b-instruct",
"meta-llama/llama-3.3-70b-instruct",
"mistralai/mistral-7b-instruct",
"meta-llama/llama-3-8b-instruct",
"allenai/olmo-7b-instruct",
"01-ai/yi-large",
"meta-llama/llama-2-70b-chat",
"meta-llama/llama-3.1-8b-instruct",
"meta-llama/llama-3.1-70b-instruct",
"meta-llama/llama-3.1-405b-instruct",
"qwen/qwen-2.5-7b-instruct",
"openai/gpt-4o",
"openchat/openchat-7b",
"ai21/jamba-instruct",
"neversleep/llama-3.1-lumimaid-8b",
"mistralai/mixtral-8x7b-instruct:nitro",
"deepseek/deepseek-r1",
"deepseek/deepseek-r1-distill-llama-70b",
"minimax/minimax-01",
"microsoft/phi-4",
"qwen/qvq-72b-preview",
]
_GPT_TURBO_MODELS = [
"gpt-3.5-turbo",
"gpt-3.5-turbo-0613",
"gpt-3.5-turbo-16k",
"gpt-3.5-turbo-16k-0613",
"gpt-3.5-turbo-1106",
"gpt-3.5-turbo-0125",
]
GPT_CHAT_MODELS = set(_GPT_4_MODELS + _GPT_TURBO_MODELS)
class OpenAIChatModel(OpenAIModel):
def _process_prompt(self, prompt: OAIChatPrompt) -> OAIChatPrompt:
return prompt
def _assert_valid_id(self, model_id: str):
if "ft:" in model_id:
model_id = model_id.split(":")[1]
assert model_id in GPT_CHAT_MODELS, f"Invalid model id: {model_id}"
@retry(stop=stop_after_attempt(8), wait=wait_fixed(2))
async def _get_dummy_response_header(self, model_id: str):
url = "https://api.openai.com/v1/chat/completions"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {openai.api_key}",
"OpenAI-Organization": self.organization,
}
data = {
"model": model_id,
"messages": [{"role": "user", "content": "Say 1"}],
}
response = requests.post(url, headers=headers, json=data)
if "x-ratelimit-limit-tokens" not in response.headers:
raise RuntimeError("Failed to get dummy response header")
return response.headers
@staticmethod
def _count_prompt_token_capacity(prompt: OAIChatPrompt, **kwargs) -> int:
# The magic formula is: .25 * (total number of characters) + (number of messages) + (max_tokens, or 15 if not specified)
BUFFER = 5 # A bit of buffer for some error margin
MIN_NUM_TOKENS = 20
num_tokens = 0
for message in prompt:
num_tokens += 1
num_tokens += len(message["content"]) / 4
return max(
MIN_NUM_TOKENS,
int(num_tokens + BUFFER)
+ kwargs.get("n", 1) * kwargs.get("max_tokens", 15),
)
def convert_top_logprobs(self, data):
# Initialize the new structure with only top_logprobs
top_logprobs = []
for item in data["content"]:
# Prepare a dictionary for top_logprobs
top_logprob_dict = {}
for top_logprob in item["top_logprobs"]:
top_logprob_dict[top_logprob["token"]] = top_logprob["logprob"]
top_logprobs.append(top_logprob_dict)
return top_logprobs
async def _make_api_call(
self, prompt: OAIChatPrompt, model_id, start_time, **params
) -> list[LLMResponse]:
LOGGER.debug(f"Making {model_id} call with {self.organization}")
if params.get("logprobs", None):
params["top_logprobs"] = params["logprobs"]
params["logprobs"] = True
api_start = time.time()
api_response: OpenAICompletion = await openai.ChatCompletion.acreate(messages=prompt, model=model_id, organization=self.organization, **params) # type: ignore
api_duration = time.time() - api_start
duration = time.time() - start_time
context_token_cost, completion_token_cost = price_per_token(model_id)
context_cost = api_response.usage.prompt_tokens * context_token_cost
completion_cost = api_response.usage.completion_tokens * completion_token_cost
return [
LLMResponse(
model_id=model_id,
completion=choice.message.content
if "tools" not in params
else choice.message.tool_calls[0]["function"]["arguments"],
stop_reason=choice.finish_reason,
api_duration=api_duration,
duration=duration,
cost=context_cost + completion_cost,
logprobs=self.convert_top_logprobs(choice.logprobs)
if choice.logprobs is not None
else None,
)
for choice in api_response.choices
]
@staticmethod
def _print_prompt_and_response(
prompts: OAIChatPrompt, responses: list[LLMResponse]
):
for prompt in prompts:
role, text = prompt["role"], prompt["content"]
cprint(f"=={role.upper()}:", "white")
cprint(text, PRINT_COLORS[role])
for i, response in enumerate(responses):
if len(responses) > 1:
cprint(f"==RESPONSE {i + 1} ({response.model_id}):", "white")
cprint(response.completion, PRINT_COLORS["assistant"], attrs=["bold"])
print()
BASE_MODELS = {
"meta-llama/Llama-3.1-8B",
"meta-llama/Llama-3.1-70B",
}
class OpenAIBaseModel(OpenAIModel):
def _process_prompt(
self, prompt: Union[OAIBasePrompt, OAIChatPrompt]
) -> OAIBasePrompt:
if isinstance(prompt, list) and isinstance(prompt[0], dict):
return messages_to_single_prompt(prompt)
return prompt
def _assert_valid_id(self, model_id: str):
assert model_id in BASE_MODELS, f"Invalid model id: {model_id}"
@retry(stop=stop_after_attempt(8), wait=wait_fixed(2))
async def _get_dummy_response_header(self, model_id: str):
url = "https://api.openai.com/v1/completions"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {openai.api_key}",
"OpenAI-Organization": self.organization,
}
data = {"model": model_id, "prompt": "a", "max_tokens": 1}
response = requests.post(url, headers=headers, json=data)
if "gpt" in model_id and "x-ratelimit-limit-tokens" not in response.headers:
raise RuntimeError("Failed to get dummy response header")
return response.headers
@staticmethod
def _count_prompt_token_capacity(prompt: OAIBasePrompt, **kwargs) -> int:
max_tokens = kwargs.get("max_tokens", 15)
n = kwargs.get("n", 1)
completion_tokens = n * max_tokens
tokenizer = tiktoken.get_encoding("cl100k_base")
if isinstance(prompt, str):
prompt_tokens = len(tokenizer.encode(prompt))
return prompt_tokens + completion_tokens
else:
prompt_tokens = sum(len(tokenizer.encode(p)) for p in prompt)
return prompt_tokens + completion_tokens
async def _make_api_call(
self, prompt: OAIBasePrompt, model_id, start_time, **params
) -> list[LLMResponse]:
LOGGER.debug(f"Making {model_id} call with {self.organization}")
api_start = time.time()
api_response: OpenAICompletion = await openai.Completion.acreate(prompt=prompt, model=model_id, organization=self.organization, **params) # type: ignore
api_duration = time.time() - api_start
duration = time.time() - start_time
if "gpt" not in model_id:
return [
LLMResponse(
model_id=model_id,
completion=choice.text,
stop_reason=choice.finish_reason,
api_duration=api_duration,
duration=duration,
cost=0,
logprobs=choice.logprobs.top_logprobs
if choice.logprobs is not None
else None,
)
for choice in api_response.choices
]
else:
context_token_cost, completion_token_cost = price_per_token(model_id)
context_cost = api_response.usage.prompt_tokens * context_token_cost
return [
LLMResponse(
model_id=model_id,
completion=choice.text,
stop_reason=choice.finish_reason,
api_duration=api_duration,
duration=duration,
cost=context_cost / len(api_response.choices)
+ count_tokens(choice.message.content) * completion_token_cost,
logprobs=choice.logprobs.top_logprobs
if choice.logprobs is not None
else None,
)
for choice in api_response.choices
]
@staticmethod
def _print_prompt_and_response(prompt: OAIBasePrompt, responses: list[LLMResponse]):
prompt_list = prompt if isinstance(prompt, list) else [prompt]
responses_per_prompt = len(responses) // len(prompt_list)
responses_list = [
responses[i : i + responses_per_prompt]
for i in range(0, len(responses), responses_per_prompt)
]
for i, (prompt, response_list) in enumerate(zip(prompt_list, responses_list)):
if len(prompt_list) > 1:
cprint(f"==PROMPT {i + 1}", "white")
if len(response_list) == 1:
cprint(f"=={response_list[0].model_id}", "white")
cprint(prompt, PRINT_COLORS["user"], end="")
cprint(
response_list[0].completion,
PRINT_COLORS["assistant"],
attrs=["bold"],
)
else:
cprint(prompt, PRINT_COLORS["user"])
for j, response in enumerate(response_list):
cprint(f"==RESPONSE {j + 1} ({response.model_id}):", "white")
cprint(
response.completion, PRINT_COLORS["assistant"], attrs=["bold"]
)
print()
# %%
core/llm_api/usage/usage_anthropic.py
+73
View File
@@ -0,0 +1,73 @@
import concurrent.futures
import os
from typing import Optional
import requests
from core.utils import setup_environment
def can_claude_api_take_n_more_concurrents(n: int) -> bool:
def ping_claude__is_ratelimited() -> Optional[bool]:
data = {
"model": "claude-2.0",
"prompt": "\n\nHuman: Count to 50.\n\nAssistant:",
"max_tokens_to_sample": 1000,
}
headers = {
"content-type": "application/json",
"accept": "application/json",
"anthropic-version": "2023-06-01",
"x-api-key": f"{os.getenv('ANTHROPIC_API_KEY')}",
}
response = requests.post(
"https://api.anthropic.com/v1/complete",
headers=headers,
json=data,
timeout=20,
)
if response.status_code == 200:
return False
elif response.status_code == 429:
return True
else:
response.raise_for_status()
# launch n threads, each of which tries to ping claude
print(f"Checking if claude can currently take {n} more concurrent requests...")
with concurrent.futures.ThreadPoolExecutor(max_workers=n) as executor:
futures = []
for _ in range(n):
futures.append(executor.submit(ping_claude__is_ratelimited))
results = [f.result() for f in futures]
result = not any(results) # if any result is true, is rate limited
print(
f"Claude currently {'can' if result else 'cannot'} take {n} more concurrent requests"
)
return result
def binary_search_for_max_concurrent_claude_requests() -> int:
min_max_new_concurrent_requests = 1
max_max_new_concurrent_requests = 100
while (
min_max_new_concurrent_requests + 5 < max_max_new_concurrent_requests
): # plus five because we don't need it accurate
mid_test_number = (
min_max_new_concurrent_requests + max_max_new_concurrent_requests
) // 2
if can_claude_api_take_n_more_concurrents(mid_test_number):
min_max_new_concurrent_requests = mid_test_number
else:
max_max_new_concurrent_requests = mid_test_number
print(
f"\n\nFinal result: Claude can currently take roughly {min_max_new_concurrent_requests} more concurrent requests.\n"
)
return min_max_new_concurrent_requests
if __name__ == "__main__":
setup_environment()
binary_search_for_max_concurrent_claude_requests()
core/llm_api/usage/usage_openai.py
+91
View File
@@ -0,0 +1,91 @@
import logging
import openai
import requests
from core.utils import setup_environment
logger = logging.getLogger(__name__)
_org_ids = {
"NYU": "org-rRALD2hkdlmLWNVCKk9PG5Xq",
"FAR": "org-AFgHGbU3MeFr5M5QFwrBET31",
"ARG": "org-4L2GWAH28buzKOIhEAb3L5aq",
}
def extract_usage(response):
requests_left = float(response.headers["x-ratelimit-remaining-requests"])
requests_limit = float(response.headers["x-ratelimit-limit-requests"])
request_usage = 1 - (requests_left / requests_limit)
tokens_left = float(response.headers["x-ratelimit-remaining-tokens"])
tokens_limit = float(response.headers["x-ratelimit-limit-tokens"])
token_usage = 1 - (tokens_left / tokens_limit)
overall_usage = max(request_usage, token_usage)
return overall_usage
def get_ratelimit_usage(data, org_id, endpoint):
try:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {openai.api_key}",
"OpenAI-Organization": org_id,
}
response = requests.post(
endpoint,
headers=headers,
json=data,
timeout=20,
)
return extract_usage(response)
except Exception as e:
logger.warning(f"Error fetching ratelimit usage: {e}")
return -1
def fetch_ratelimit_usage(org_id, model_name) -> float:
data = {
"model": model_name,
"messages": [{"role": "user", "content": "Say 1"}],
}
return get_ratelimit_usage(
data, org_id, "https://api.openai.com/v1/chat/completions"
)
def fetch_ratelimit_usage_base(org_id, model_name) -> float:
data = {"model": model_name, "prompt": "a", "max_tokens": 1}
return get_ratelimit_usage(data, org_id, "https://api.openai.com/v1/completions")
def get_current_openai_model_usage() -> None:
models_to_check = [
"gpt-3.5-turbo-instruct",
"gpt-4-1106-preview",
"gpt-4-base",
]
org_names = ["NYU", "FAR", "ARG"]
result_str = (
"\nModel usage: 1 is hitting rate limits, 0 is not in use. -1 is error.\n"
)
for org in org_names:
result_str += f"\n{org}:\n"
for model_name in models_to_check:
if model_name == "gpt-4-base" or model_name == "gpt-3.5-turbo-instruct":
if org == "ARG":
usage = fetch_ratelimit_usage_base(_org_ids[org], model_name)
else:
continue
else:
usage = fetch_ratelimit_usage(_org_ids[org], model_name)
result_str += f"\t{model_name}:\t{usage:.2f}\n"
result_str += "\n"
print(result_str)
if __name__ == "__main__":
setup_environment()
get_current_openai_model_usage()
core/utils.py
+239
View File
@@ -0,0 +1,239 @@
import asyncio
import json
import logging
import os
import time
from functools import lru_cache, wraps
from pathlib import Path
from typing import Callable
import matplotlib.pyplot as plt
import numpy as np
import openai
import pandas as pd
import replicate
import typer
import yaml
from tenacity import retry, retry_if_result, stop_after_attempt
typer.main.get_command_name = lambda name: name
LOGGER = logging.getLogger(__name__)
SEPARATOR = "---------------------------------------------\n\n"
SEPARATOR_CONVERSATIONAL_TURNS = "=============================================\n\n"
PROMPT_HISTORY = "prompt_history"
SECRETS_FILE_PATH = Path(__file__).parent.parent / "SECRETS"
LOGGING_LEVELS = {
"critical": logging.CRITICAL,
"error": logging.ERROR,
"warning": logging.WARNING,
"info": logging.INFO,
"debug": logging.DEBUG,
}
def setup_environment(
anthropic_tag: str = "ANTHROPIC_API_KEY",
logger_level: str = "info",
openai_tag: str = "API_KEY",
mistral_tag: str = "MISTRAL_API_KEY",
replicate_tag: str = "REPLICATE_API_KEY",
organization: str = None,
):
setup_logging(logger_level)
load_secrets(
SECRETS_FILE_PATH,
anthropic_tag,
logger_level,
openai_tag,
mistral_tag,
replicate_tag,
organization,
)
def setup_logging(level_str):
level = LOGGING_LEVELS.get(
level_str.lower(), logging.INFO
) # default to INFO if level_str is not found
logging.basicConfig(
level=level,
format="%(asctime)s [%(levelname)s] (%(name)s) %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
root_logger = logging.getLogger()
root_logger.setLevel(level)
# Disable logging from noisy libraries
logging.getLogger("openai").setLevel(logging.CRITICAL)
logging.getLogger("httpx").setLevel(logging.CRITICAL)
logging.getLogger("matplotlib").setLevel(logging.CRITICAL)
logging.getLogger("anthropic").setLevel(logging.CRITICAL)
logging.getLogger("httpcore").setLevel(logging.CRITICAL)
logging.getLogger("urllib3").setLevel(logging.CRITICAL)
LOGGER.info(f"Logging level set to {level_str}")
def load_secrets(
file_path=SECRETS_FILE_PATH,
anthropic_tag: str = "ANTHROPIC_API_KEY",
logger_level: str = "info",
openai_tag: str = "API_KEY",
mistral_tag: str = "MISTRAL_API_KEY",
replicate_tag: str = "REPLICATE_API_KEY",
organization: str = None,
):
secrets = {}
with open(file_path) as f:
for line in f:
key, value = line.strip().split("=", 1)
secrets[key] = value
openai.api_key = secrets[openai_tag]
os.environ['LLAMA_API_BASE'] = secrets['LLAMA_API_BASE']
# replicate.api_token = secrets[replicate_tag]
# os.environ["ANTHROPIC_API_KEY"] = secrets[anthropic_tag]
# os.environ["MISTRAL_API_KEY"] = secrets[mistral_tag]
# os.environ["REPLICATE_API_KEY"] = secrets[replicate_tag]
if organization is not None:
openai.organization = secrets[organization]
if secrets.get("API_BASE") is not None:
openai.api_base = secrets['API_BASE']
return secrets
def load_yaml(file_path):
with open(file_path) as f:
content = yaml.safe_load(f)
return content
@lru_cache(maxsize=8)
def load_yaml_cached(file_path):
with open(file_path) as f:
content = yaml.safe_load(f)
return content
def save_yaml(file_path, data):
with open(file_path, "w") as f:
yaml.dump(data, f)
def load_jsonl(file_path):
data = []
with open(file_path, "r") as f:
for line in f:
json_obj = json.loads(line)
data.append(json_obj)
return data
def save_jsonl(file_path, data):
with open(file_path, "w") as f:
for line in data:
json.dump(line, f)
f.write("\n")
def delete_old_prompt_files(
path: str = PROMPT_HISTORY, max_age_minutes: int = 60, keep_recent: int = 50
):
"""
Delete all files in the folder that:
- Are more than max_age_minutes old
- AND are not one of the keep_recent most recent files
"""
if not os.path.exists(path):
return
# Get all files in the folder with their full paths and creation times
files = [
{
"path": os.path.join(path, filename),
"ctime": os.path.getctime(os.path.join(path, filename)),
}
for filename in os.listdir(path)
if os.path.isfile(os.path.join(path, filename))
]
# Sort files by creation time
files.sort(key=lambda f: f["ctime"], reverse=True)
# Current time in seconds since epoch
now = time.time()
deleted_count = 0
for index, file_info in enumerate(files):
# File age in minutes
age_minutes = (now - file_info["ctime"]) / 60
# If file is older than x_minutes and is not one of the y_most_recent files, delete it
if age_minutes > max_age_minutes and index >= keep_recent:
os.remove(file_info["path"])
deleted_count += 1
if deleted_count > 0:
print(f"Deleted {deleted_count} old prompt files")
def typer_async(f):
@wraps(f)
def wrapper(*args, **kwargs):
try:
loop = asyncio.get_running_loop()
except RuntimeError: # No event loop running
loop = None
if loop is None:
return asyncio.run(f(*args, **kwargs))
else:
return f(*args, **kwargs) # Return coroutine to be awaited
return wrapper
@retry(
stop=stop_after_attempt(16),
retry=retry_if_result(lambda result: result is not True),
)
def function_with_retry(function, *args, **kwargs):
return function(*args, **kwargs)
@retry(
stop=stop_after_attempt(16),
retry=retry_if_result(lambda result: result is not True),
)
async def async_function_with_retry(function, *args, **kwargs):
return await function(*args, **kwargs)
def log_model_timings(api_handler, save_location="./model_timings.png"):
if len(api_handler.model_timings) > 0:
plt.figure(figsize=(10, 6))
for model in api_handler.model_timings:
timings = np.array(api_handler.model_timings[model])
wait_times = np.array(api_handler.model_wait_times[model])
LOGGER.info(
f"{model}: response {timings.mean():.3f}, waiting {wait_times.mean():.3f} (max {wait_times.max():.3f}, min {wait_times.min():.3f})"
)
plt.plot(
timings, label=f"{model} - Response Time", linestyle="-", linewidth=2
)
plt.plot(
wait_times, label=f"{model} - Waiting Time", linestyle="--", linewidth=2
)
plt.legend()
plt.title("Model Performance: Response and Waiting Times")
plt.xlabel("Sample Number")
plt.ylabel("Time (seconds)")
plt.savefig(save_location, dpi=300)
plt.close()
def softmax(x):
return np.exp(x) / np.sum(np.exp(x), axis=0)
env.yaml
+176
View File
@@ -0,0 +1,176 @@
name: UE
channels:
- conda-forge
- defaults
dependencies:
- _libgcc_mutex=0.1=main
- _openmp_mutex=5.1=1_gnu
- bzip2=1.0.8=h5eee18b_6
- ca-certificates=2024.7.2=h06a4308_0
- expat=2.6.2=h6a678d5_0
- ld_impl_linux-64=2.38=h1181459_1
- libffi=3.4.4=h6a678d5_1
- libgcc-ng=11.2.0=h1234567_1
- libgomp=11.2.0=h1234567_1
- libstdcxx-ng=11.2.0=h1234567_1
- libuuid=1.41.5=h5eee18b_0
- ncurses=6.4=h6a678d5_0
- openssl=3.0.14=h5eee18b_0
- pip=24.2=py312h06a4308_0
- python=3.12.4=h5148396_1
- readline=8.2=h5eee18b_0
- setuptools=72.1.0=py312h06a4308_0
- shortuuid=1.0.13=pyhd8ed1ab_0
- sqlite=3.45.3=h5eee18b_0
- tk=8.6.14=h39e8969_0
- wheel=0.43.0=py312h06a4308_0
- xz=5.4.6=h5eee18b_1
- zlib=1.2.13=h5eee18b_1
- pip:
- adjusttext==1.2.0
- aiohappyeyeballs==2.4.3
- aiohttp==3.10.10
- aiosignal==1.3.1
- alembic==1.13.3
- altair==5.4.1
- annotated-types==0.7.0
- anthropic==0.37.1
- antlr4-python3-runtime==4.9.3
- anyio==4.6.2.post1
- attrs==24.2.0
- beautifulsoup4==4.12.3
- blobfile==3.0.0
- bs4==0.0.2
- cattrs==24.1.2
- certifi==2024.8.30
- cfgv==3.4.0
- charset-normalizer==3.4.0
- click==8.1.7
- contourpy==1.3.0
- cycler==0.12.1
- datasets==3.0.2
- dill==0.3.8
- distlib==0.3.9
- distro==1.9.0
- docker-pycreds==0.4.0
- eval-type-backport==0.2.0
- fastapi==0.100.0
- filelock==3.16.1
- fire==0.7.0
- fonttools==4.54.1
- frozenlist==1.4.1
- fsspec==2024.9.0
- gitdb==4.0.11
- gitpython==3.1.43
- greenlet==3.1.1
- h11==0.14.0
- httpcore==1.0.6
- httptools==0.6.4
- httpx==0.27.2
- huggingface-hub==0.26.1
- hydra-core==1.3.2
- identify==2.6.1
- idna==3.10
- jinja2==3.1.4
- jiter==0.6.1
- joblib==1.4.2
- jsonpath-python==1.0.6
- jsonschema==4.23.0
- jsonschema-specifications==2024.10.1
- kiwisolver==1.4.7
- lxml==5.3.0
- mako==1.3.6
- markdown-it-py==3.0.0
- markupsafe==3.0.2
- matplotlib==3.9.2
- mdurl==0.1.2
- mistralai==1.1.0
- mpmath==1.3.0
- multidict==6.1.0
- multiprocess==0.70.16
- mypy-extensions==1.0.0
- narwhals==1.10.0
- networkx==3.4.2
- nodeenv==1.9.1
- numpy==2.1.2
- nvidia-cublas-cu12==12.4.5.8
- nvidia-cuda-cupti-cu12==12.4.127
- nvidia-cuda-nvrtc-cu12==12.4.127
- nvidia-cuda-runtime-cu12==12.4.127
- nvidia-cudnn-cu12==9.1.0.70
- nvidia-cufft-cu12==11.2.1.3
- nvidia-curand-cu12==10.3.5.147
- nvidia-cusolver-cu12==11.6.1.9
- nvidia-cusparse-cu12==12.3.1.170
- nvidia-nccl-cu12==2.21.5
- nvidia-nvjitlink-cu12==12.4.127
- nvidia-nvtx-cu12==12.4.127
- omegaconf==2.3.0
- openai==0.28.0
- packaging==24.1
- pandas==2.2.3
- pebble==5.0.7
- pillow==10.4.0
- platformdirs==4.3.6
- polars==1.10.0
- pre-commit==4.0.1
- propcache==0.2.0
- protobuf==5.28.2
- psutil==6.1.0
- pyarrow==17.0.0
- pycryptodomex==3.21.0
- pydantic==2.9.2
- pydantic-core==2.23.4
- pygments==2.18.0
- pyparsing==3.2.0
- python-dateutil==2.8.2
- python-dotenv==1.0.1
- pytz==2024.2
- pyyaml==6.0.2
- referencing==0.35.1
- regex==2024.9.11
- replicate==1.0.2
- requests==2.32.3
- rich==13.9.3
- rpds-py==0.20.0
- safetensors==0.4.5
- scikit-learn==1.5.2
- scipy==1.14.1
- seaborn==0.13.2
- sentry-sdk==2.17.0
- setproctitle==1.3.3
- shellingham==1.5.4
- six==1.16.0
- smmap==5.0.1
- sniffio==1.3.1
- soupsieve==2.6
- sqlalchemy==2.0.18
- starlette==0.27.0
- sympy==1.13.1
- tabulate==0.9.0
- tenacity==9.0.0
- termcolor==2.5.0
- threadpoolctl==3.5.0
- tiktoken==0.8.0
- together==1.3.4
- tokenizers==0.20.1
- torch==2.5.1
- tqdm==4.66.5
- transformers==4.46.3
- triton==3.1.0
- trueskill==0.4.5
- typer==0.12.5
- typing-extensions==4.12.2
- typing-inspect==0.9.0
- tzdata==2024.2
- urllib3==2.2.3
- uvicorn==0.22.0
- uvloop==0.21.0
- virtualenv==20.27.0
- wandb==0.18.5
- watchfiles==0.24.0
- websockets==13.1
- xxhash==3.5.0
- yarl==1.16.0
- zstandard==0.23.0
prefix: /home/wenjiaxin/miniconda3/envs/UE
figures/algorithm.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 131 KiB

figures/claude_performance.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 191 KiB

figures/llama_performance.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 232 KiB

requirements.txt
+31
View File
@@ -0,0 +1,31 @@
adjustText
alembic
altair
anthropic
cattrs
datasets
fastapi==0.100.0
fire
hydra-core
matplotlib
mistralai
openai==0.28.0
pandas
pebble
polars
pre-commit
pydantic
replicate
scikit-learn
scipy
seaborn
sqlalchemy==2.0.18
tabulate
tenacity
termcolor
tiktoken
tqdm
trueskill
typer[all]
uvicorn[standard]==0.22.0
wandb
setup.py
+11
View File
@@ -0,0 +1,11 @@
# setup.py
from setuptools import find_packages, setup
setup(
name="my_package",
version="0.1",
packages=find_packages(),
install_requires=[
# List your package dependencies here
],
)
src/datatypes/enums.py
+34
View File
@@ -0,0 +1,34 @@
__all__ = ["Language", "PromptType"]
from enum import Enum
class Language(Enum):
PYTHON = ("python", "Python")
CPP = ("cpp", "C++")
def __init__(self, code, text):
self.code = code
self.text = text
@staticmethod
def from_code(code):
if code == "python":
return Language.PYTHON
elif code == "cpp":
return Language.CPP
else:
raise Exception(f"Unknown code langauge: {code}")
class PromptType(Enum):
SOLUTION = "solution_generation"
BLUE_TEAM = "blue_team"
RED_TEAM = "red_team"
EVAL = "eval"
class DifficultyEstimationType(Enum):
PROBLEM_ONLY = "problem_only"
PROBLEM_SOLUTION = "problem_solution"
PROBLEM_SOLUTION_EXECUTION = "problem_solution_execution"
src/experiments/ICM.py
+560
View File
@@ -0,0 +1,560 @@
import asyncio
import json
import math
import os
import random
from collections import Counter
from copy import deepcopy
from tqdm import tqdm
import numpy as np
from datasets import load_dataset
import argparse
from core.llm_api.llm import ModelAPI
from core.utils import setup_environment
from src.experiments.ICM_tools import (
propose_consistencyfix,
run_consistencyfix,
pick_two_inconsistent_claims,
update_assign_based_on_decision,
)
from src.model_querying.prompt_creation import (
get_decision_prompt,
get_judge_prompt_fewshot,
)
from src.model_querying.solution_extraction import (
extract_claim_logprobs,
extract_decision_logprobs,
)
from src.pipeline.pipeline import Pipeline, PipelineConfig
from src.tools.dataloaders import (
load_assignments,
load_problems_from_json,
load_problems_from_json_ids,
)
from src.tools.path_utils import get_default_results_directory, get_root_directory
def calculate_accuracy(train_data, inconsistent_pairs):
train_probs = []
for i in train_data.values():
if i["label"] is None:
continue
if i["label"] == 1:
train_probs.append(i["score"])
else:
train_probs.append(-i["score"])
if len(train_probs) == 0:
train_prob = 0
else:
train_prob = np.mean(train_probs)
return {
"train_accuracy": 0
if len(train_data) == 0
else np.mean([i["label"] == i["vanilla_label"] for i in train_data.values()]),
"train_label_distribution": Counter(
[i["vanilla_label"] for i in train_data.values()]
),
"train_predict_distribution": Counter(
[i["label"] for i in train_data.values()]
),
"train_prob": train_prob,
"train_size": len(train_data),
"inconsistent_num": len(inconsistent_pairs),
}
def update_assign(data):
for key, value in data.items():
if value["score"] > 0:
value["label"] = 1
else:
value["label"] = 0
return data
def fix_inconsistency(demonstrations, cur_metric, name, alpha, iter=0, K=20):
backup_metric = deepcopy(cur_metric)
if cur_metric["inconsistent_num"] == 0:
return demonstrations, cur_metric
cur_pool = {k: v for k, v in demonstrations.items() if v["label"] is not None}
assignment = cur_pool
best_metric = cur_metric
best_assignment = assignment
best_decision_id = None
for k in range(K):
pipeline = propose_consistencyfix(
args.model,
name=name,
iter=f"{iter}-{k}",
assignment=assignment,
)
results = asyncio.run(pipeline.run())
decisions = results["decisions"]
assignment = results["get_assign"]
for decision_id, decision in enumerate(decisions.values()):
tmp_decision_metric_list = []
tmp_decision_assignment_list = []
for score_idx, score in enumerate([0, 1]):
tmp_decision = deepcopy(decision)
tmp_decision["score"] = score
tmp_assignment = update_assign_based_on_decision(
deepcopy(assignment), tmp_decision
)
tmp_pipeline = run_consistencyfix(
model=args.model,
name=name,
iter=f"{iter}-{k}-{decision_id}-{score_idx}",
assignment=tmp_assignment,
)
tmp_results = asyncio.run(tmp_pipeline.run())
tmp_metric = tmp_results["evaluate"]
tmp_decision_metric_list.append(tmp_metric)
tmp_decision_assignment_list.append(tmp_assignment)
tmp_best_decision_id = np.argmax(
[get_energy(i, args.alpha) for i in tmp_decision_metric_list]
)
tmp_assignment = tmp_decision_assignment_list[tmp_best_decision_id]
tmp_metric = tmp_decision_metric_list[tmp_best_decision_id]
if get_energy(tmp_metric, args.alpha) >= get_energy(best_metric, args.alpha):
best_decision_id = decision_id
best_metric = tmp_metric
best_assignment = tmp_assignment
break
if best_decision_id is None:
break
elif best_metric["inconsistent_num"] == 0:
assignment = best_assignment
break
else:
assignment = best_assignment
for k in assignment:
demonstrations[k] = assignment[k]
return demonstrations, best_metric
def get_pipeline(
model,
name=None,
use_cache=True,
num_problems=None,
decision_id=None,
iter=None,
assignment=None,
):
pipeline_name = f"iterative-truth-assign-iter-{iter}"
if decision_id is not None:
pipeline_name += f"-{decision_id}"
if name is not None:
pipeline_name += "-" + name
ROOT_DIR = get_root_directory()
DATA_DIR = ROOT_DIR / "data"
pipeline_config = PipelineConfig(
pipeline_name,
anthropic_num_threads=40,
openai_fraction_rate_limit=0.99,
num_problems=num_problems,
use_cache=use_cache,
)
pipeline = Pipeline(pipeline_config)
assert assignment is not None
initial_assign = pipeline.add_load_data_step(
"get_assign", load_assignments, assignment
)
def add_train_demonstrations(train_data):
copy_data = deepcopy(train_data)
copy_data = {k: v for k, v in copy_data.items() if v["label"] is not None}
keys = list(copy_data.keys())
values = list(copy_data.values())
saved_keys = [
"prompt",
"question",
"choice",
"choice_2",
"consistency_id",
"consistency_key",
"source",
"label",
"vanilla_label",
]
values = []
for i in copy_data.values():
values.append({saved_key: i[saved_key] for saved_key in saved_keys if saved_key in i})
for idx, key in enumerate(keys):
tmp_keys, tmp_values = [], []
for j, (prev_key, prev_value) in enumerate(zip(keys, values)):
if j != idx:
tmp_keys.append(prev_key)
tmp_values.append(prev_value)
demos = {
prev_key: prev_value
for j, (prev_key, prev_value) in enumerate(zip(tmp_keys, tmp_values))
}
sorted_demos = {}
for k, v in demos.items():
q = v["consistency_id"]
if q not in sorted_demos:
sorted_demos[q] = []
sorted_demos[q].append((k, v))
out_sorted_demos = {}
for group in sorted_demos.values():
for k, v in group:
out_sorted_demos[k] = v
copy_data[key]["demonstration"] = out_sorted_demos
return copy_data
merged_train_data = pipeline.add_transformation_step(
"add_train_demonstration",
add_train_demonstrations,
dependencies=[initial_assign],
)
get_train_preds = pipeline.add_query_step(
"get_train_preds",
model,
get_judge_prompt_fewshot,
extract_claim_logprobs,
dependencies=[merged_train_data],
logprobs=20,
max_tokens=1,
use_cache=use_cache,
)
pick_claims = pipeline.add_transformation_step(
"pick_two_inconsistent_claims",
pick_two_inconsistent_claims,
dependencies=[initial_assign],
)
eval_preds = pipeline.add_eval_step(
"evaluate",
calculate_accuracy,
dependencies=[get_train_preds, pick_claims],
)
return pipeline
async def predict_assignment(model, example, demonstrations):
demos = [
v
for k, v in demonstrations.items()
if k != example["uid"] and v["label"] is not None
]
anthropic_requests = [
model_api(
model,
get_judge_prompt_fewshot(
example,
demos,
pipeline=False,
),
logprobs=20,
max_tokens=1,
parse_fn=extract_claim_logprobs,
)
]
responses = await asyncio.gather(*anthropic_requests)
score = responses[0][0]["score"]
new_label = score > 0
return int(new_label)
def get_temperature(
iteration, initial_temp, final_temp, decay_rate, schedule="exp"
):
"""
Calculate the temperature for simulated annealing.
Parameters:
- iteration: Current iteration number.
- initial_temp: Initial temperature.
- decay_rate: Rate at which the temperature decreases.
Returns:
- Current temperature.
"""
if schedule == "exp":
return max(final_temp, initial_temp * (decay_rate**iteration))
elif schedule == "log":
return max(final_temp, initial_temp / (1 + 2 * np.log(1 + iteration)))
else:
assert False
def get_energy(metric, alpha):
return alpha * metric["train_prob"] - metric["inconsistent_num"]
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--alpha", type=float, default=30)
parser.add_argument("--seed", type=int, default=27565976)
parser.add_argument("--testbed", type=str, default="gsm8k")
parser.add_argument("--model", type=str, default="meta-llama/Llama-3.1-70B")
parser.add_argument("--batch_size", type=int, default=256)
parser.add_argument("--num_seed", type=int, default=8)
parser.add_argument("--K", type=int, default=3000)
parser.add_argument("--consistency_fix_K", type=int, default=10)
parser.add_argument("--decay", type=float, default=0.99)
parser.add_argument("--initial_T", type=float, default=10)
parser.add_argument("--final_T", type=float, default=0.01)
parser.add_argument("--scheduler", type=str, default="log")
args = parser.parse_args()
return args
def load_data(args):
if args.testbed == "alpaca":
with open(get_root_directory() / "data/train_alpaca.json") as f:
train = json.load(f)
template = """Human: {question}
Response 1: {choice}
Response 2: {choice_2}
Claim: Response 1 is more helpful and harmless than Response 2.
I think this claim is """
for i in train:
i['source'] = 'alpaca'
i['consistency_key'] = 'A>B' if i['label'] else 'B>A'
i['prompt'] = template.format(question=i['question'], choice=i['choice'].replace("\n", ""), choice_2=i['choice_2'].replace("\n", ""))
args.GROUP_SIZE = 2
elif args.testbed == 'gsm8k':
with open(get_root_directory() / "data/train_gsm8k.json") as f:
train = json.load(f)
template = """Question: {question}
Claim: {answer}
I think this claim is """
for i in train:
i['source'] = 'gsm8k'
i['consistency_key'] = i['answer']
i['prompt'] = template.format(question=i['question'], answer=i['choice'])
args.GROUP_SIZE = 4
elif args.testbed == 'truthfulQA':
with open(get_root_directory() / "data/train_truthfulqa.json") as f:
train = json.load(f)
template = """Question: {question}
Claim: {answer}
I think this claim is """
for i in train:
i['source'] = 'truthfulQA'
i['consistency_key'] = 'A' if i['label'] else 'B'
i['prompt'] = template.format(question=i['question'], answer=i['choice'])
args.GROUP_SIZE = 4
elif args.testbed == 'truthfulQA-preference':
with open(get_root_directory() / "data/train_truthfulqa_preference.json") as f:
train = json.load(f)
template = """Question: {question}
Answer 1: {choice}
Answer 2: {choice_2}
Claim: Answer 1 is more truthful than Answer 2.
I think this claim is """
for i in train:
i['source'] = 'truthfulQA-preference'
i['consistency_key'] = 'A>B' if i['label'] else 'B>A'
i['prompt'] = template.format(question=i['question'], choice=i['choice'], choice_2=i['choice_2'])
args.GROUP_SIZE = 2
train_map = {}
for i in train:
if i['consistency_id'] not in train_map:
train_map[i['consistency_id']] = []
train_map[i['consistency_id']].append(i)
out = []
for key in train_map:
out += train_map[key]
train = out
# sample a batch of batch_size datapoints
fewshot_ids = random.sample(
list(range(len(train)// args.GROUP_SIZE)), args.batch_size // args.GROUP_SIZE
)
fewshot_ids = [
i * args.GROUP_SIZE + j for i in fewshot_ids for j in range(args.GROUP_SIZE)
]
return train, fewshot_ids
def initialize(train, fewshot_ids, args):
demonstrations = {}
unlabeled_ids = []
whole_ids = []
seed_ids = []
random_init_labels = [1] * (args.num_seed // 2) + [0] * (args.num_seed // 2)
random.shuffle(random_init_labels)
for id, i in enumerate(fewshot_ids):
item = train[i]
item["vanilla_label"] = item["label"] # store dataset labels to measure agreement during the searching process
item["uid"] = id
whole_ids.append(item["uid"])
if id >= args.num_seed: # set labels to None
item["label"] = None
item["type"] = "predict"
unlabeled_ids.append(item["uid"])
else: # set random labels
item["type"] = "seed"
item["label"] = random_init_labels[id]
seed_ids.append(item["uid"])
demonstrations[id] = item
return demonstrations, unlabeled_ids, whole_ids, seed_ids
def main(args):
train, fewshot_ids = load_data(args)
demonstrations, unlabeled_ids, whole_ids, seed_ids = initialize(train, fewshot_ids, args)
cur_metric = {
"train_prob": -1e6,
"inconsistent_num": 100000,
"train_accuracy": 1.0,
"train_predict_distribution": {"0": 0, "1": 0},
"train_label_distribution": {"0": 0, "1": 0},
}
print('init random labels = ', Counter([i['label'] for i in demonstrations.values() if i['type'] == 'seed']), 'init label acc = ', np.mean([i['label'] == i['vanilla_label'] for i in demonstrations.values() if i['type'] == 'seed']))
name = f"{args.testbed}-llama70b-K{args.K}-bc{args.batch_size}_seed{args.seed}-initialsize{args.num_seed}-weighted{args.alpha}-decay{args.decay}-initialT{args.initial_T}-finalT{args.final_T}-scheduler{args.scheduler}"
iter = 0
flip_cnt = 0
example_id = 0
for _ in tqdm(range(args.K), desc="searching"):
cur_pool = {
k: v for k, v in demonstrations.items() if v["label"] is not None
}
initial_demos = deepcopy(demonstrations)
if iter == 0:
pipeline = get_pipeline(
args.model,
name=name,
num_problems=None,
iter=iter,
assignment=cur_pool,
)
results = asyncio.run(pipeline.run())
cur_metric = results["evaluate"]
demonstrations, cur_metric = fix_inconsistency(
demonstrations, cur_metric, name, args.alpha, iter=iter, K=args.consistency_fix_K
)
cur_pool = {
k: v for k, v in demonstrations.items() if v["label"] is not None
}
while True: # weighted sampling
candidates_ids = whole_ids
weights = [1 for _ in range(len(candidates_ids))]
for i in candidates_ids:
if i in cur_pool:
same_consistency_group_ids = [j for j in candidates_ids if demonstrations[j]["consistency_id"] == demonstrations[i]["consistency_id"]]
for j in same_consistency_group_ids:
if j not in cur_pool:
weights[j] = 100
example_id = random.choices(candidates_ids, k=1, weights=weights)[0]
break
new_label = asyncio.run(
predict_assignment(
args.model,
demonstrations[example_id],
cur_pool,
)
)
if demonstrations[example_id]["label"] != new_label:
tmp_demonstrations = deepcopy(demonstrations)
tmp_demonstrations[example_id]["label"] = new_label
dummy_metric = {
"train_prob": -1e6,
"inconsistent_num": 100000,
"train_accuracy": 1.0,
"train_predict_distribution": {"0": 0, "1": 0},
"train_label_distribution": {"0": 0, "1": 0},
}
tmp_demonstrations, _ = fix_inconsistency(
tmp_demonstrations,
dummy_metric,
name + "newlabelexplore",
args.alpha,
iter=iter,
K=10,
)
tmp_pool = {
k: v
for k, v in tmp_demonstrations.items()
if v["label"] is not None
}
pipeline = get_pipeline(
model=args.model,
name=name,
num_problems=None,
iter=iter,
assignment=tmp_pool,
)
results = asyncio.run(pipeline.run())
metric = results["evaluate"]
T = get_temperature(
flip_cnt, args.initial_T, args.final_T, args.decay, schedule=args.scheduler
)
print(f"iter = {iter}, pool size = {len(cur_pool)}, cur acc = {cur_metric['train_accuracy']}, new acc = {metric['train_accuracy']}, cur score = {get_energy(cur_metric, args.alpha)}, new score = {get_energy(metric, args.alpha)}, cur inconsistent num = {cur_metric['inconsistent_num']}, new inconsistent num = {metric['inconsistent_num']}")
print('cur label distribution = ', Counter([i['label'] for i in demonstrations.values() if i['label'] is not None]))
print('new label distribution = ', Counter([i['label'] for i in tmp_demonstrations.values() if i['label'] is not None]))
accept_prob = math.exp((get_energy(metric, args.alpha) - get_energy(cur_metric, args.alpha)) / T)
print("accept prob = ", accept_prob)
if random.random() < accept_prob:
print("accept")
demonstrations = tmp_demonstrations
flip_cnt += 1
cur_metric = metric
with open(f"log_{name}.jsonl", "a") as f:
f.write(json.dumps({
"iter": iter,
"flip_cnt": flip_cnt,
"acc": cur_metric['train_accuracy'],
"score": get_energy(cur_metric, args.alpha),
}) + "\n")
else:
print("reject")
print("=" * 100)
iter += 1
if __name__ == "__main__":
setup_environment(logger_level="error")
model_api = ModelAPI(anthropic_num_threads=20, openai_fraction_rate_limit=0.99)
args = get_args()
print("task: ", args.testbed)
random.seed(args.seed)
main(args)
src/experiments/ICM_tools.py
+238
View File
@@ -0,0 +1,238 @@
import asyncio
import json
import random
from collections import Counter
from copy import deepcopy
import numpy as np
from datasets import load_dataset
from src.model_querying.prompt_creation import (
get_decision_prompt,
get_judge_prompt_fewshot,
)
from src.model_querying.solution_extraction import (
extract_claim_logprobs,
extract_decision_logprobs,
)
from src.pipeline.pipeline import Pipeline, PipelineConfig
from src.tools.dataloaders import (
load_assignments,
load_problems_from_json,
load_problems_from_json_ids,
)
from src.tools.path_utils import get_default_results_directory, get_root_directory
def calculate_accuracy(train_data, inconsistent_pairs):
return {
"train_predict_distribution": Counter(
[i["label"] for i in train_data.values()]
),
"train_label_distribution": Counter(
[i["vanilla_label"] for i in train_data.values()]
),
"train_accuracy": np.mean(
[i["label"] == i["vanilla_label"] for i in train_data.values()]
),
"train_prob": np.mean(
[
i["score"] if i["label"] == 1 else -i["score"]
for i in train_data.values()
]
),
"train_size": len(train_data),
"inconsistent_num": len(inconsistent_pairs),
}
def update_assign_based_on_decision(data, decision):
if decision["type"] == "contradiction":
if decision["score"] > 0:
data[decision["claim_1"]["uid"]]["label"] = 1
data[decision["claim_2"]["uid"]]["label"] = 0
else:
data[decision["claim_1"]["uid"]]["label"] = 0
data[decision["claim_2"]["uid"]]["label"] = 1
else:
assert decision["type"] == "implication"
if decision["score"] > 0:
data[decision["claim_1"]["uid"]]["label"] = 1
data[decision["claim_2"]["uid"]]["label"] = 1
else:
data[decision["claim_1"]["uid"]]["label"] = 0
data[decision["claim_2"]["uid"]]["label"] = 0
return data
def pick_two_inconsistent_claims(data):
claims = list(data.values())
consistency_groups = {}
for claim in claims:
cid = claim["consistency_id"]
if cid not in consistency_groups:
consistency_groups[cid] = []
consistency_groups[cid].append(claim)
inconsistent_pairs = {}
for group in consistency_groups.values():
labels = [claim["vanilla_label"] for claim in group]
for i in range(len(group)):
for j in range(i + 1, len(group)):
if (group[i]['consistency_key'] != group[j]['consistency_key']) and (
(group[i]['label'] == group[j]['label'] == 1) or
(
(group[i]['consistency_key'] in ['A>B', 'B>A']) and (group[i]['label'] == group[j]['label'] == 0) # in comparative tasks, at least one of the two claims is correct
)
):
# if (group[i]["vanilla_label"] != group[j]["vanilla_label"]) and (
# group[i]["label"] == group[j]["label"]
# ):
inconsistent_pairs[len(inconsistent_pairs)] = {
"claim_1": group[i],
"claim_2": group[j],
"consistency_id": group[i]["consistency_id"],
"type": "contradiction",
}
elif (group[i]["consistency_key"] == group[j]["consistency_key"]) and (
group[i]["label"] != group[j]["label"]
):
inconsistent_pairs[len(inconsistent_pairs)] = {
"claim_1": group[i],
"claim_2": group[j],
"consistency_id": group[i]["consistency_id"],
"type": "implication",
}
random.shuffle(inconsistent_pairs)
return inconsistent_pairs
def propose_consistencyfix(
model,
name=None,
iter=None,
assignment=None,
use_cache=True,
):
pipeline_name = f"propose-consistencyfix-iter-{iter}"
if name is not None:
pipeline_name += "-" + name
pipeline_config = PipelineConfig(
pipeline_name,
anthropic_num_threads=40,
openai_fraction_rate_limit=0.99,
num_problems=None,
use_cache=use_cache,
)
pipeline = Pipeline(pipeline_config)
initial_assign = pipeline.add_load_data_step(
"get_assign", load_assignments, assignment
)
pick_claims = pipeline.add_transformation_step(
"pick_two_inconsistent_claims",
pick_two_inconsistent_claims,
dependencies=[initial_assign],
)
get_decision = pipeline.add_query_step(
"decisions",
model,
get_decision_prompt,
extract_decision_logprobs,
dependencies=[pick_claims],
logprobs=20,
max_tokens=1,
use_cache=use_cache,
)
return pipeline
def run_consistencyfix(
model,
name=None,
use_cache=True,
decision_id=None,
decision=None,
iter=None,
assignment=None,
):
pipeline_name = f"consistencyfix-iter-{iter}"
if decision_id is not None:
pipeline_name += f"-{decision_id}"
if name is not None:
pipeline_name += "-" + name
pipeline_config = PipelineConfig(
pipeline_name,
anthropic_num_threads=40,
openai_fraction_rate_limit=0.99,
num_problems=None,
use_cache=use_cache,
)
pipeline = Pipeline(pipeline_config)
assert assignment is not None
initial_assign = pipeline.add_load_data_step(
"get_assign", load_assignments, assignment
)
pick_claims = pipeline.add_transformation_step(
"pick_two_inconsistent_claims",
pick_two_inconsistent_claims,
dependencies=[initial_assign],
)
def add_train_demonstrations(train_data):
copy_data = deepcopy(train_data)
keys = list(copy_data.keys())
values = list(copy_data.values())
saved_keys = [
"prompt",
"question",
"choice",
"choice_2",
"consistency_id",
"source",
"label",
"vanilla_label",
]
values = []
for i in copy_data.values():
values.append(
{saved_key: i[saved_key] for saved_key in saved_keys if saved_key in i}
)
for idx, key in enumerate(keys):
train_data[key]["demonstration"] = {
prev_key: prev_value
for j, (prev_key, prev_value) in enumerate(zip(keys, values))
if j != idx
}
return train_data
merged_train_data = pipeline.add_transformation_step(
"add_train_demonstration",
add_train_demonstrations,
dependencies=[initial_assign],
)
get_train_preds = pipeline.add_query_step(
"get_train_preds",
model,
get_judge_prompt_fewshot,
extract_claim_logprobs,
dependencies=[merged_train_data],
logprobs=20,
max_tokens=1,
use_cache=use_cache,
)
eval_preds = pipeline.add_eval_step(
"evaluate",
calculate_accuracy,
dependencies=[get_train_preds, pick_claims],
)
return pipeline
src/experiments/log_truthfulQA-llama70b-K3000-bc256_seed27565976-initialsize8-weighted50.0-decay0.99-initialT10-finalT0.01-schedulerlog.jsonl
+6
View File
@@ -0,0 +1,6 @@
{"iter": 2, "flip_cnt": 1, "acc": 0.7777777777777778, "score": 59.95838513901126}
{"iter": 4, "flip_cnt": 2, "acc": 0.8, "score": 62.33137680262533}
{"iter": 7, "flip_cnt": 3, "acc": 0.9090909090909091, "score": 71.59853190532465}
{"iter": 8, "flip_cnt": 4, "acc": 0.9166666666666666, "score": 79.43930320079279}
{"iter": 10, "flip_cnt": 5, "acc": 0.9230769230769231, "score": 82.68683171642135}
{"iter": 11, "flip_cnt": 6, "acc": 0.9285714285714286, "score": 86.60106027413124}
src/experiments/plot.py
+24
View File
@@ -0,0 +1,24 @@
import json
from matplotlib import pyplot as plt
import os
for file in os.listdir("."):
if file.startswith("log_"):
print(file)
with open(file, "r") as f:
data = [json.loads(i) for i in f]
data = data[:60]
x = list(range(8, 8 + len(data)))
y_score = [i['score'] for i in data]
y_acc = [i['acc'] for i in data]
# plt.subplot(1, 2, 1)
plt.plot(x, y_score, label=f'Acc {max(y_acc):.2f}')
plt.xlabel("# Searched Claims")
plt.ylabel("Score")
# plt.subplot(1, 2, 2)
# plt.plot(x, y_acc)
# plt.xlabel("# Searched Claims")
# plt.ylabel("Accuracy")
plt.legend()
plt.tight_layout()
plt.savefig("log.png", dpi=500)
src/model_querying/__init__.py
View File
src/model_querying/prompt_creation.py
+45
View File
@@ -0,0 +1,45 @@
__all__ = ["get_apps_solution_prompt", "get_two_shot_prompt", "get_backdoor_prompt"]
import json
import random
from src.datatypes.enums import Language, PromptType
from src.tools.dataloaders import load_prompts
from src.tools.string_manipulation import format_key_suffix, strip_comments_from_string
SOLUTION_PROMPTS = load_prompts(PromptType.SOLUTION)
RED_TEAM_PROMPTS = load_prompts(PromptType.RED_TEAM)
BLUE_TEAM_PROMPTS = load_prompts(PromptType.BLUE_TEAM)
EVAL_PROMPTS = load_prompts(PromptType.EVAL)
class Prompt:
@staticmethod
def empty():
return Prompt([])
def __init__(self, text, logit_bias=None):
self.text = text
self.logit_bias = logit_bias
def get_decision_prompt(example):
prompt = 'dummy prompt'
return Prompt(prompt)
def get_judge_prompt_fewshot(example, demonstrations=None, pipeline=True):
if demonstrations is None:
demonstrations = list(example["demonstration"].values())
prompt = ""
for i in demonstrations:
prompt += i['prompt']
prompt += "True" if i["label"] else "False"
prompt += "\n\n"
prompt += example['prompt']
if pipeline:
return Prompt(prompt)
else:
return prompt
src/model_querying/solution_extraction.py
+55
View File
@@ -0,0 +1,55 @@
import json
import logging
import math
from copy import copy
logger = logging.getLogger(__name__)
def get_yes_no(x):
x = x.lower()
y = "true" in x
n = "false" in x
if y == n:
return None
return y
def get_yes_no_diff_logprobs(logprobs):
eps = 1e-5
prob_sums = {False: eps, True: eps}
for k, v in logprobs.items():
o = get_yes_no(k)
if o is None:
continue
prob_sums[o] += math.exp(v)
if prob_sums[False] == eps and prob_sums[True] == eps:
return 0
else:
return math.log(prob_sums[True]) - math.log(prob_sums[False])
def extract_claim_logprobs(response):
response = response.copy()
try:
logprobs = response["response"]["logprobs"][0]
response[f"score"] = get_yes_no_diff_logprobs(logprobs)
except Exception as e:
logger.info(
f"Problem {response['metadata']['uid']}: Error extracting judgment: {repr(e)}"
)
response["score"] = 0
return response
def extract_decision_logprobs(response):
response = response.copy()
try:
logprobs = response["response"]["logprobs"][0]
response[f"score"] = get_yes_no_diff_logprobs(logprobs)
except Exception as e:
logger.info(
f"Problem {response['metadata']['uid']}: Error extracting decision: {repr(e)}"
)
response["score"] = 0
return response
src/pipeline/README.md
+27
View File
@@ -0,0 +1,27 @@
How to use the pipeline:
First, outline the graph you would like to execute, including the following:
* Data Loading
* Model queries
* Code Execution Eval
* Transformations
* Monitoring
Next, convert each of the nodes in that graph into the corresponding helper function:
* add_load_data_step
* add_query_step
* add_code_evaluation_step
* add_transformation_step
* add_monitoring_step
Each of these takes different parameters that you can see in the method signatures. The important ones to know are these:
LoadData takes either a data-loading function and a location, or it takes a dataset
Queries take a prompt function that they pass the incoming data into to create the associated prompt, and a parse function that they use to parse the LLM response
Code Evals take an executor function that executes all of the code in the Solution objects on the associated test cases.
Transforms take arbitrary functions that they apply to the data as a whole.
Monitoring steps take arbitrary monitoring steps. I may eventually enforce that all pipelines end in one of these because it's really what we care about.
Finally, put the dependencies of each step into their dependencies parameter. This is how execution order is determined and how data flows between steps. If you rely on more than one step, the data will be passed to ordered args in the same order as the list of dependencies.
You will also need to include a PipelineConfig parameter that contains metadata around how many concurrents to use and similar.
Once you have a pipeline definition, call the pipeline.run() function on it to execute the graph. This method returns the Pipeline.Results object back, which holds the output of each step in a dictionary.
src/pipeline/pipeline.py
+314
View File
@@ -0,0 +1,314 @@
__all__ = ["PipelineConfig", "Pipeline"]
import asyncio
import logging
from collections import deque
from tqdm.auto import tqdm
from core.llm_api.llm import ModelAPI
from src.datatypes.enums import Language
from src.runners.query_model import QueryConfigBuilder, query_model
from src.tools.dataloaders import read_from_cache, save_to_cache
from src.tools.path_utils import get_root_directory
logger = logging.getLogger(__name__)
def in_notebook():
try:
from IPython import get_ipython
if get_ipython() is None or "IPKernelApp" not in get_ipython().config:
return False
except ImportError:
return False
return True
class Task:
def __init__(self, name, func, use_cache, dependencies=[]):
self.name = name
self.func = func
self.use_cache = use_cache
self.index = None
self.dependencies = dependencies
self.dependents = []
self.result = None
for dep in dependencies:
dep.dependents.append(self)
async def execute(self, results):
if self.result is None:
dep_results = [results[dep.name] for dep in self.dependencies]
if asyncio.iscoroutinefunction(self.func):
self.result = await self.func(
*dep_results, use_cache=self.use_cache, index=self.index
)
else:
self.result = self.func(
*dep_results, use_cache=self.use_cache, index=self.index
)
return self.result
class PipelineConfig:
def __init__(
self,
name,
anthropic_num_threads=2,
openai_fraction_rate_limit=0.99,
use_cache=True,
language=Language.PYTHON,
num_problems=None,
problem_ids=None,
num_open_files=1000000,
organization="NYU_ORG",
print_prompt_and_response=False,
api_base=None,
):
self.name = name
self.anthropic_num_threads = anthropic_num_threads
self.openai_fraction_rate_limit = openai_fraction_rate_limit
self.organization = organization
self.print_prompt_and_response = print_prompt_and_response
self.use_cache = use_cache
self.language = language
self.num_problems = num_problems
self.problem_ids = problem_ids
self.num_open_files = num_open_files
self.api_base = api_base
self.play_sound = in_notebook()
class Pipeline:
def __init__(self, config):
self.config = config
self.steps = []
self.step_names = set()
self.results = {}
self.model_api = ModelAPI(
self.config.anthropic_num_threads,
self.config.openai_fraction_rate_limit,
self.config.organization,
self.config.print_prompt_and_response,
)
self.file_sem = asyncio.BoundedSemaphore(self.config.num_open_files)
self.cost = {"red": 0, "blue": 0}
def add_load_data_step(
self, name, dataloader_fn, data_location, dependencies=[], use_cache=None
):
if name in self.step_names:
raise ValueError(f"Step name {name} already exists")
self.step_names.add(name)
def call(*args, use_cache, index):
return dataloader_fn(
data_location,
num_problems=self.config.num_problems,
problem_ids=self.config.problem_ids,
)
task = Task(name, call, use_cache, dependencies)
self.steps.append(task)
return task
def add_query_step(
self,
name,
model,
prompt_fn,
parse_fn,
dependencies=[],
use_cache=None,
temperature=None,
logprobs=None,
team=None,
max_tokens=4096,
bon=1,
):
if name in self.step_names:
raise ValueError(f"Step name {name} already exists")
self.step_names.add(name)
query_config_builder = (
QueryConfigBuilder()
.with_model_to_test(model)
.with_prompt_fn(lambda x: prompt_fn(x))
.with_parse_fn(lambda x: parse_fn(x))
.with_num_problems(self.config.num_problems)
.with_max_tokens(max_tokens)
.with_temperature(temperature)
.with_logprobs(logprobs)
.with_bon(bon)
)
async def call(data, use_cache, index):
response_dict = await query_model(
self.model_api,
self.file_sem,
query_config_builder.with_experiment_name(
f"{self.config.name}/{index:02d}-{name}"
)
.with_use_cache(use_cache)
.with_data(data)
.build(),
)
self.add_cost_data(team, response_dict)
return response_dict
step = Task(name, call, use_cache, dependencies)
self.steps.append(step)
return step
def add_transformation_step(
self,
name,
transformation_fn,
dependencies=[],
use_cache=None,
strong_model=None,
weak_model=None,
read_cache=False,
):
if name in self.step_names:
raise ValueError(f"Step name {name} already exists")
self.step_names.add(name)
async def call(*args, use_cache, index):
incoming_problem_ids = set().union(*[arg.keys() for arg in args])
if use_cache and read_cache:
logger.debug(
f"Reading from cache for transformation: {self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}"
)
data, cached_problem_ids = read_from_cache(
f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}"
)
if incoming_problem_ids.issubset(set(cached_problem_ids)):
return {k: v for k, v in data.items() if k in incoming_problem_ids}
if asyncio.iscoroutinefunction(transformation_fn):
output = await transformation_fn(*args)
else:
output = transformation_fn(*args)
async with self.file_sem:
save_to_cache(
output,
f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}",
delete_existing=read_cache,
incoming_problem_ids=incoming_problem_ids,
)
return output
step = Task(name, call, use_cache, dependencies)
self.steps.append(step)
return step
def add_eval_step(
self,
name,
eval_fn,
dependencies=[],
strong_model=None,
weak_model=None,
):
if name in self.step_names:
raise ValueError(f"Step name {name} already exists")
self.step_names.add(name)
async def call(*args, use_cache, index):
output = eval_fn(*args)
cache_obj = {"summary": output}
async with self.file_sem:
save_to_cache(
cache_obj,
f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}",
)
return output
step = Task(name, call, None, dependencies)
self.steps.append(step)
return step
def topological_sort_tasks(self, tasks):
in_degree = {task: len(task.dependencies) for task in tasks}
queue = deque([task for task in tasks if in_degree[task] == 0])
sorted_tasks = []
task_order = {task: i for i, task in enumerate(tasks)}
while queue:
task = queue.popleft()
sorted_tasks.append(task)
for dependent in task.dependents:
in_degree[dependent] -= 1
if in_degree[dependent] == 0:
queue.append(dependent)
queue = deque(sorted(queue, key=lambda t: task_order[t]))
for i, task in enumerate(sorted_tasks):
task.index = i
return sorted_tasks
def add_cost_data(self, team, response_dict):
cost = sum(
[response["response"]["cost"] for response in response_dict.values()]
)
if team is not None:
if team not in self.cost:
self.cost[team] = 0
self.cost[team] += cost
overall_team = team.split("_")[0]
if overall_team != team:
self.cost[overall_team] += cost
def set_use_cache(self, tasks):
# This is called after tasks.sort, so we are guaranteed to process all
# dependencies before each task itself.
for task in tasks:
if not self.config.use_cache:
task.use_cache = False
continue
if task.use_cache is None:
task.use_cache = True
for dep in task.dependencies:
if not dep.use_cache:
task.use_cache = False
def speak(self, message):
if self.config.play_sound:
from IPython.display import Javascript, display
display(
Javascript(
f"""
if(window.speechSynthesis) {{
var synth = window.speechSynthesis;
synth.speak(new window.SpeechSynthesisUtterance('{message}'));
}}
"""
)
)
async def run(self):
steps = self.topological_sort_tasks(self.steps)
self.set_use_cache(steps)
for task in steps:
logger.info(
f"Starting step {task.index}: {task.name} - Using cache: {task.use_cache}"
)
try:
self.results[task.name] = await task.execute(self.results)
except Exception as e:
logger.error(f"Error in step {task.index}: {task.name}")
logger.error(e)
self.speak("Pipeline failed sad face")
raise e
logger.info(f"Finished step {task.index}: {task.name}")
self.speak("Jobs done")
logger.info("Run complete!! Nice!! 🚀🚀")
return self.results
src/runners/__init__.py
View File
src/runners/evaluate_code.py
+191
View File
@@ -0,0 +1,191 @@
__all__ = [
"EvalConfig",
"EvalConfigBuilder",
"evaluate_solutions",
"examine_solution",
"print_eval",
]
import json
import os
from src.code_evaluation.test_results import Solution
import src.tools.path_utils as path_utils
DEFAULT_RESULTS_DIR = path_utils.get_default_results_directory()
class EvalConfig:
def __init__(
self,
experiment_name,
model_to_test,
executor_fn,
language,
use_cache=True,
data=None,
dataloader_fn=None,
data_location=None,
):
self.experiment_name = experiment_name
self.model_to_test = model_to_test
if dataloader_fn is not None:
self.data = dataloader_fn(data_location)
elif isinstance(data, str):
with open(json.load(data), "r") as f:
self.data = json.load(f)
else:
self.data = data
self.executor_fn = executor_fn
self.language = language
self.use_cache = use_cache
def __str__(self):
pass
def __repr__(self):
return self.__str__()
class EvalConfigBuilder:
def __init__(self):
self.experiment_name = None
self.model_to_test = None
self.executor_fn = None
self.language = None
self.data = None
self.use_cache = None
self.dataloader_fn = None
self.data_location = None
def with_experiment_name(self, experiment_name):
self.experiment_name = experiment_name
return self
def with_executor_fn(self, executor_fn):
self.executor_fn = executor_fn
return self
def with_language(self, language):
self.language = language
return self
def with_model_to_test(self, model_to_test):
self.model_to_test = model_to_test
return self
def with_use_cache(self, use_cache):
self.use_cache = use_cache
return self
def with_data(self, data):
self.data = data
return self
def with_dataloader_fn(self, dataloader_fn):
self.dataloader_fn = dataloader_fn
return self
def with_data_location(self, data_location):
self.data_location = data_location
return self
def build(self):
assert self.experiment_name is not None, "Experiment name must be set"
assert self.executor_fn is not None, "Executor function must be set"
assert self.language is not None, "Language must be set"
assert (self.data is not None) or (
self.dataloader_fn is not None and self.data_location is not None
), "Data must be set"
return EvalConfig(
self.experiment_name,
self.model_to_test,
self.executor_fn,
self.language,
self.use_cache,
self.data,
self.dataloader_fn,
self.data_location,
)
def evaluate_solutions(eval_config):
if isinstance(eval_config.data, list):
eval_config.data = {
index: response[0] for index, response in enumerate(eval_config.data)
}
eval_data = []
for problem_id, response in eval_config.data.items():
if response == {}:
item = Solution.no_solution(problem_id)
else:
item = Solution.from_response(problem_id, response, eval_config.language)
eval_data.append(item)
results_dir = DEFAULT_RESULTS_DIR
save_dir = results_dir / eval_config.experiment_name / eval_config.model_to_test
os.makedirs(save_dir, exist_ok=True)
# load caching results
if eval_config.use_cache and save_dir is not None:
for idx, solution in enumerate(eval_data):
if solution.correct is not None: # already have runtime eval results
continue
save_path = f"{save_dir}/{solution.question_id}.json"
if os.path.exists(save_path):
eval_data[idx] = Solution.from_cache_file(save_path)
executor_results = eval_config.executor_fn(eval_data)
for problem_id, problem_data in eval_config.data.items():
if problem_id not in executor_results:
executor_results[problem_id] = {}
for field, value in problem_data.items():
if field not in executor_results[problem_id]:
executor_results[problem_id][field] = value
# save caching results
if save_dir is not None:
for question_id, solution in executor_results.items():
save_path = f"{save_dir}/{question_id}.json"
with open(save_path, "w") as f:
json.dump(solution, f, indent=2)
return executor_results
def examine_solution(solutions, index):
print(
f"Difficulty:\n{solutions[index][0]['metadata']['difficulty']}\n----------------------------"
)
print(solutions[index][0]["metadata"]["question"])
print(solutions[index][0]["solution"])
print("test cases:")
for test in solutions[index][0]["metadata"]["test_cases"]:
print(f"{test['input']}{test['output']}")
def print_eval(results):
correct_tests = 0
total_tests = 0
correct_problems = 0
total_problems = len(results.keys())
for problem_id, problem_result in results.items():
correct_tests_local = sum(
[
1
for test_result in problem_result["test_cases"]
if test_result["correct"]
]
)
total_tests_local = len(problem_result["test_cases"])
# print(f"Problem ID: {problem_id}\nTest Results:\n\tCorrect: {correct_tests_local}\n\tTotal: {total_tests_local}\n\tAccuracy: {(correct_tests_local * 100.)/total_tests_local}\nOverall Correct: {problem_result.correct}")
correct_tests += correct_tests_local
total_tests += total_tests_local
if problem_result["correct"]:
correct_problems += 1
print(
f"Number of Problems: {total_problems}\nNumber Correct: {correct_problems}\nAccuracy: {(correct_problems * 100.)/total_problems}\nNumber of Tests: {total_tests}\nNumber Correct: {correct_tests}\nAccuracy: {(correct_tests * 100.)/total_tests}"
)
src/runners/query_model.py
+276
View File
@@ -0,0 +1,276 @@
__all__ = ["QueryConfig", "QueryConfigBuilder", "query_model"]
import asyncio
import os
import tiktoken
import src.tools.path_utils as path_utils
ROOT_DIR = path_utils.get_root_directory()
DEFAULT_RESULTS_DIR = path_utils.get_default_results_directory()
class QueryConfig:
def __init__(
self,
experiment_name,
model_to_test,
dataloader_fn,
data_location,
data,
prompt_fn,
parse_fn=None,
use_cache=False,
num_problems=None,
max_tokens=4096,
results_dir=None,
temperature=None,
logprobs=None,
bon=1,
):
assert isinstance(model_to_test, str)
self.experiment_name = experiment_name
self.model_to_test = model_to_test
self.dataloader_fn = dataloader_fn
self.data_location = data_location
self.data = data
self.prompt_fn = prompt_fn
self.use_cache = use_cache
self.parse_fn = parse_fn
self.num_problems = num_problems
self.max_tokens = max_tokens
self.results_dir = results_dir
self.temperature = temperature if temperature is not None else 0.0
self.logprobs = logprobs
self.bon = bon
def get_data(self):
assert self.data is not None
return self.data
def __str__(self):
return (
f"QueryConfig("
f"experiment_name={self.experiment_name}, "
f"model_to_test={self.model_to_test}, "
f"dataloader_fn={self.dataloader_fn}, "
f"data_location={self.data_location}, "
f"data={self.data}, "
f"prompt_fn={self.prompt_fn}, "
f"use_cache={self.use_cache}, "
f"num_problems={self.num_problems}, "
f"max_tokens={self.max_tokens}, "
f"results_dir={self.results_dir}, "
f"temperature={self.temperature}, "
f"logprobs={self.logprobs}"
)
def __repr__(self):
return self.__str__()
class QueryConfigBuilder:
def __init__(self):
self.experiment_name = None
self.model_to_test = None
self.dataloader_fn = None
self.data_location = None
self.data = None
self.prompt_fn = None
self.parse_fn = None
self.use_cache = False
self.num_problems = None
self.max_tokens = 4096
self.results_dir = None
self.temperature = 0.0
self.logprobs = None
self.bon = 1
def with_experiment_name(self, experiment_name):
self.experiment_name = experiment_name
return self
def with_bon(self, bon):
self.bon = bon
return self
def with_model_to_test(self, model_to_test):
assert isinstance(model_to_test, str)
self.model_to_test = model_to_test
return self
def with_dataloader_fn(self, dataloader_fn):
self.dataloader_fn = dataloader_fn
return self
def with_data_location(self, data_location):
self.data_location = data_location
return self
def with_data(self, data):
self.data = data
return self
def with_prompt_fn(self, prompt_fn):
self.prompt_fn = prompt_fn
return self
def with_parse_fn(self, parse_fn):
self.parse_fn = parse_fn
return self
def with_use_cache(self, use_cache):
self.use_cache = use_cache
return self
def with_num_problems(self, num_problems):
self.num_problems = num_problems
return self
def with_max_tokens(self, max_tokens):
self.max_tokens = max_tokens
return self
def with_results_dir(self, results_dir):
self.results_dir = results_dir
return self
def with_temperature(self, temperature):
self.temperature = temperature
return self
def with_logprobs(self, logprobs):
self.logprobs = logprobs
if logprobs is not None:
assert "claude" not in self.model_to_test
return self
def build(self):
assert self.experiment_name is not None, "Experiment name must be set"
assert self.model_to_test is not None, "Model to test must be set"
assert self.prompt_fn is not None, "Prompt function must be set"
assert (self.data is not None) or (
self.dataloader_fn is not None and self.data_location is not None
), "Data must be set"
assert (self.data is None) or (
self.dataloader_fn is None and self.data_location is None
), "Data and dataloader_fn/data_location cannot both be set"
return QueryConfig(
self.experiment_name,
self.model_to_test,
self.dataloader_fn,
self.data_location,
self.data,
self.prompt_fn,
self.parse_fn,
self.use_cache,
self.num_problems,
self.max_tokens,
self.results_dir,
self.temperature,
self.logprobs,
self.bon,
)
def _get_prompts(problems, prompt_fn):
prompts = {}
for problem_id, problem in problems.items():
prompt = prompt_fn(problem)
if prompt.text:
prompts[problem_id] = prompt
return prompts
def get_save_dir(query_config):
results_dir = (
ROOT_DIR / query_config.results_dir
if query_config.results_dir is not None
else DEFAULT_RESULTS_DIR
)
save_dir = results_dir / query_config.experiment_name / query_config.model_to_test
os.makedirs(save_dir, exist_ok=True)
return save_dir
def move_data_into_metadata(data):
for data_id, value in data.items():
filtered_val = {
k: v
for k, v in value.items()
if k not in ("metadata", "prompt", "response")
}
metadata = value.get("metadata", {})
new_metadata = metadata | filtered_val
data[data_id]["metadata"] = new_metadata
def format_response(data, model_responses_map, query_config):
model_responses_flattened = {}
for data_id, response in model_responses_map.items():
if query_config.bon == 1:
model_responses_flattened[f"{data_id}"] = data[data_id] | response[0]
continue
for resp_id, resp in enumerate(response):
model_responses_flattened[f"{data_id}-{resp_id}"] = data[data_id] | resp
return model_responses_flattened
def tokenize_logit_bias(logit_bias, model):
tokenizer = tiktoken.encoding_for_model(model)
tokenized_bias = {}
for k, v in logit_bias.items():
tokenized = tokenizer.encode(k)
assert len(tokenized) == 1, f"Tokenized bias key {k} is not a single token"
tokenized_bias[tokenized[0]] = v
return tokenized_bias
async def query_model(model_api, file_sem, query_config):
data = query_config.get_data()
prompts = _get_prompts(data, query_config.prompt_fn)
save_dir = get_save_dir(query_config)
move_data_into_metadata(data)
model_requests = [
model_api(
query_config.model_to_test,
prompts[data_id].text,
max_tokens=query_config.max_tokens,
temperature=query_config.temperature,
n=query_config.bon,
top_p=1.0,
logprobs=query_config.logprobs,
use_cache=query_config.use_cache,
metadata=data[data_id]["metadata"],
parse_fn=query_config.parse_fn,
save_path=f"{save_dir}/{data_id}.json",
file_sem=file_sem,
**(
{
"logit_bias": tokenize_logit_bias(
prompts[data_id].logit_bias, query_config.model_to_test
)
}
if prompts[data_id].logit_bias is not None
else {}
),
)
for data_id in prompts.keys()
]
model_responses = await asyncio.gather(*model_requests)
# pass through data that wasn't modified by the request
model_responses_map = {
data_id: response for data_id, response in zip(prompts.keys(), model_responses)
}
for key in data.keys():
if key not in prompts:
model_responses_map[key] = [data[key]]
response = format_response(data, model_responses_map, query_config)
return response
src/tools/__init__.py
View File
src/tools/dataloaders.py
+264
View File
@@ -0,0 +1,264 @@
__all__ = ["load_prompts", "load_problems", "load_problems_from_json", "load_solutions"]
import json
import logging
import os
from .path_utils import get_default_results_directory, get_root_directory
logger = logging.getLogger(__name__)
ROOT_DIR = get_root_directory()
DATA_DIR = ROOT_DIR / "data" / "APPS"
PROMPTS_DIR = ROOT_DIR / "src" / "prompts"
def get_data_dir():
return DATA_DIR
def load_prompts(prompt_type):
files = [f for f in (PROMPTS_DIR / prompt_type.value).glob("*") if f.is_file()]
prompts = {}
for file in files:
with file.open("r") as f:
prompts[file.name] = f.read()
return prompts
def load_problem_subset(subset, require_solutions=False, problem_ids=None):
def load_problems(dir, num_problems=None):
if problem_ids:
problem_dirs = problem_ids
else:
problem_dirs = os.listdir(dir)
problems = {}
added = 0
for problem_dir in problem_dirs:
problem_path = dir / problem_dir
problem = {}
with (problem_path / "metadata.json").open("r") as f:
problem["metadata"] = json.load(f)
if (
subset != "ALL"
and problem["metadata"]["difficulty"].lower() != subset.lower()
):
continue
if require_solutions and not (problem_path / "solutions.json").exists():
logger.debug(
f"Skipping problem {problem_dir} because it does not have solutions"
)
continue
with (problem_path / "question.txt").open("r") as f:
problem["question"] = f.read()
problem["uid"] = problem_dir
problems[problem_dir] = problem
added += 1
if added >= num_problems:
break
return problems
return load_problems
def load_problems(dir, num_problems=None):
return load_problem_subset("ALL")(dir, num_problems)
def load_problems_from_json(path, num_problems=None, problem_ids=None):
problems = {}
try:
with open(path) as f:
data = json.load(f)
except Exception as e:
print('read data error: ', e)
data = path
if num_problems is not None:
data = data[:num_problems]
for i, item in enumerate(data):
item["uid"] = i
if 'vanilla_label' not in item:
item["vanilla_label"] = item["label"]
problems[f"{i}"] = item
return problems
def load_problems_from_json_ids(path, num_problems=None, problem_ids=None):
problems = {}
try:
with open(path) as f:
data = json.load(f)
except:
data = path
if problem_ids is not None:
data = [data[i] for i in problem_ids]
for i, item in enumerate(data):
item["uid"] = i
if 'vanilla_label' not in item:
item["vanilla_label"] = item["label"]
# acc.append(item['label'] == item['vanilla_label'])
problems[f"{i}"] = item
return problems
def load_assignments(path, num_problems=None, problem_ids=None):
return path
def load_solutions(dir, num_problems=None):
solutions = {}
files = dir.glob("*")
files = [f for f in files if f.name != "incoming_problem_ids.json"]
if num_problems is not None:
files = list(files)[:num_problems]
for file in files:
with file.open("r") as f:
solution = json.load(f)
# Assume 1 solution per file
if isinstance(solution, list):
solutions[file.stem] = solution[0]
else:
solutions[file.stem] = solution
return solutions
def load_multiple_solutions(dir, num_problems=None, problem_ids=None):
solutions = {}
files = dir.glob("*")
files = [f for f in files if f.name != "incoming_problem_ids.json"]
for file in files:
with file.open("r") as f:
solution = json.load(f)
if "metadata" in solution:
solution.pop("metadata")
if "demonstration" in solution:
solution.pop("demonstration")
solutions[file.stem] = solution
return solutions
def load_multiple_solutions_w2s(dir, num_problems=None, problem_ids=None):
solutions = {}
files = dir.glob("*")
files = [f for f in files if f.name != "incoming_problem_ids.json"]
for file in files:
with file.open("r") as f:
solution = json.load(f)
metadata = solution[0]['metadata']
if "demonstration" in metadata:
metadata.pop("demonstration")
metadata['label'] = solution[0]['score'] > 0
solutions[file.stem] = metadata
return solutions
def save_to_cache(data, name, delete_existing=False, incoming_problem_ids=None):
dir = get_default_results_directory() / name
# Delete all files in the directory first
if delete_existing and os.path.exists(dir):
for file in os.listdir(dir):
file_path = os.path.join(dir, file)
if os.path.isfile(file_path):
os.unlink(file_path)
os.makedirs(dir, exist_ok=True)
for k, v in data.items():
if isinstance(v, list):
to_write = [
{
key: value
for key, value in item.items()
if key not in ["prompt", "response"]
}
for item in v
]
else:
to_write = {
key: value
for key, value in v.items()
if key not in ["prompt", "response"]
}
with open(dir / f"{k}.json", "w") as f:
json.dump(to_write, f, indent=4)
if incoming_problem_ids:
with open(dir / "incoming_problem_ids.json", "w") as f:
json.dump({"problem_ids": list(incoming_problem_ids)}, f, indent=4)
def read_from_cache(name):
dir = get_default_results_directory() / name
data = {}
incoming_problem_ids = []
for file in dir.glob("*.json"):
if file.name == "incoming_problem_ids.json":
with file.open("r") as f:
incoming_problem_ids = json.load(f).get("problem_ids", [])
else:
with file.open("r") as f:
value = json.load(f)
if not value.get("metadata"):
value["metadata"] = {k: v for k, v in value.items()}
data[file.stem] = value
return data, incoming_problem_ids
def load_ground_truth_solutions(problem_ids):
output = {}
for problem_id in problem_ids:
with open(get_data_dir() / "test" / problem_id / "solutions.json", "r") as f:
solutions = json.load(f)
cleaned_solutions = []
for solution in solutions:
# Remove unwanted lines from the solution
cleaned_solution = []
for line in solution.split("\n"):
if (
not line.strip().startswith("#!")
and " input=" not in line
and "sys.stdin" not in line
):
cleaned_solution.append(line)
cleaned_solutions.append("\n".join(cleaned_solution).strip())
output[problem_id] = cleaned_solutions
return output
def load_test_case(problem_id):
problem_id = problem_id.split("-")[0]
with open(get_data_dir() / "test" / problem_id / "input_output.json", "r") as f:
data = json.load(f)
return [
{"input": i, "output": o} for (i, o) in zip(data["inputs"], data["outputs"])
]
def load_test_cases(problem_ids):
output = {}
for problem_id in problem_ids:
output[problem_id] = load_test_case(problem_id)
return output
loaded_test_cases = {}
def get_test_cases_for_single_problem(problem_id):
global loaded_test_cases
if problem_id not in loaded_test_cases:
loaded_test_cases[problem_id] = load_test_case(problem_id)
return loaded_test_cases[problem_id]
src/tools/path_utils.py
+11
View File
@@ -0,0 +1,11 @@
__all__ = ["get_root_directory", "get_default_results_directory"]
from pathlib import Path
def get_root_directory():
return Path(__file__).parent.parent.parent
def get_default_results_directory():
return get_root_directory() / "results"
src/tools/printer.py
+134
View File
@@ -0,0 +1,134 @@
import json
import os
import subprocess
from pathlib import Path
from anytree import Node, RenderTree
from anytree.exporter import DotExporter
from src.tools.path_utils import get_default_results_directory
def print_experiment_log(experiment_name, strong_model, weak_model, problem_number):
results_dir = get_default_results_directory()
experiment_dir = results_dir / experiment_name
# Get all step directories
step_dirs = [d for d in experiment_dir.iterdir() if d.is_dir()]
# Sort step directories by step number and exclude "merged_results"
step_dirs = [d for d in step_dirs if d.name != "merged_results"]
step_dirs.sort(key=lambda x: int(x.name.split("-")[0]))
for step_dir in step_dirs:
# Check for both strong and weak model directories
for model in [strong_model, weak_model, f"{strong_model}+{weak_model}"]:
model_dir = step_dir / model
if not model_dir.exists():
continue
ignore_keys = ["metadata", "prompt", "response"]
if model == f"{strong_model}+{weak_model}":
ignore_keys.extend(["question", "test_cases", "uid"])
# Find matching problem files
problem_files = list(model_dir.glob(f"{problem_number}*.json"))
for problem_file in problem_files:
with open(problem_file, "r") as f:
data = json.load(f)
print(f"Step: {step_dir.name}")
print(f"Model: {model}")
print(f"Problem: {problem_file.stem}")
if not isinstance(data, list):
data = [data]
print("\nPrompt:")
prompt_array = data[0].get("prompt")
if prompt_array is None:
print("No prompt available")
else:
for text in prompt_array:
print(f"Role: {text['role']}")
print(f"Content: {text['content']}")
for response in data:
print("\nResponse:")
print(
response.get("response", {}).get(
"completion", "No response available"
)
)
print("\nOther Fields:")
for key, value in response.items():
if key not in ignore_keys:
print(f"{key}: {value}")
print("\n" + "=" * 50 + "\n")
def show_pipeline_graph(pipeline):
import matplotlib.pyplot as plt
import networkx as nx
# Create a directed graph
G = nx.DiGraph()
# Add nodes and edges
for task in pipeline.steps:
G.add_node(task.name)
for dep in task.dependencies:
G.add_edge(dep.name, task.name)
# Print the graph structure
print("Pipeline Dependency Graph:")
for node in nx.topological_sort(G):
predecessors = list(G.predecessors(node))
successors = list(G.successors(node))
print(f"{node}:")
if predecessors:
print(f" Parents: {', '.join(predecessors)}")
if successors:
print(f" Children: {', '.join(successors)}")
# Generate a DOT file for visualization
output_dir = get_default_results_directory() / pipeline.config.name
output_dir.mkdir(parents=True, exist_ok=True)
dot_file = output_dir / "pipeline_graph.dot"
png_file = output_dir / "pipeline_graph.png"
nx.drawing.nx_pydot.write_dot(G, str(dot_file))
print(f"DOT file generated at: {dot_file}")
# Generate PNG file using Graphviz
try:
subprocess.run(["dot", "-Tpng", str(dot_file), "-o", str(png_file)], check=True)
print(f"PNG file generated at: {png_file}")
except subprocess.CalledProcessError:
print(
"Error: Failed to generate PNG. Make sure Graphviz is installed and accessible in your PATH."
)
except FileNotFoundError:
print(
"Error: Graphviz not found. Please install Graphviz to generate PNG files."
)
# Optionally, you can also use matplotlib to visualize the graph
plt.figure(figsize=(12, 8))
pos = nx.spring_layout(G)
nx.draw(
G,
pos,
with_labels=True,
node_color="lightblue",
node_size=2000,
font_size=8,
arrows=True,
)
plt.title("Pipeline Dependency Graph")
plt.axis("off")
plt.tight_layout()
plt.savefig(str(output_dir / "pipeline_graph_matplotlib.png"))
print(
f"Matplotlib graph generated at: {output_dir / 'pipeline_graph_matplotlib.png'}"
)
src/tools/string_manipulation.py
+62
View File
@@ -0,0 +1,62 @@
import re
def format_key_suffix(key_suffix):
if key_suffix:
if key_suffix[0] != "_":
key_suffix = f"_{key_suffix}"
else:
key_suffix = ""
return key_suffix
COMMENTS_REGEX = r"""
^ # Begin of line.
(?:
# A) Capturing group n°1: Full-line comment followed by empty lines.
(
[ \t]* # Optional spaces or tabs.
\#[^\r\n]*\r?\n # The comment and the new line.
(?:[ \t]*\r?\n)* # Optional empty lines (perhaps with spaces/tabs).
)
|
# B) Statement and optional comment at the end.
(?:
( # Capturing group n°2 : The statement
(?:
# Multi-line strings with \"\"\" or '''.
# Capturing group n°3 : The triple quotes.
(['\"]{3})[\s\S]*?\3
|
# Double-quoted string "It's ok".
\"(?: \\. | [^\"] )*\"
|
# Single-quoted string 'I\'ll say "Hello!"'.
'(?: \\. | [^'] )*'
|
# Any chars, except spaces, hashtag, quotes and new lines.
[^ \t#\"'\r\n]+
|
# Horizontal spaces, but not followed by a comment, because
# we want the spaces in front of the comment to be matched
# together with the optional comment we want to get rid of.
[ \t]+(?![ \t]*\#)
)+
)
# Capturing group n°4: An optional comment at the end of a statement.
(
[ \t]*\#[^\r\n]*
)?
)+
)
"""
def strip_comments_from_string(string):
return re.sub(
COMMENTS_REGEX,
"\\2",
string,
0,
re.MULTILINE | re.VERBOSE | re.UNICODE,
)