stanford-oval--storm
shaoyijia
507 行
19 KiB
Python
507 行
19 KiB
Python
import logging
|
|
import os
|
|
import random
|
|
import threading
|
|
from typing import Optional, Literal, Any
|
|
|
|
import backoff
|
|
import dspy
|
|
import requests
|
|
from dsp import ERRORS, backoff_hdlr, giveup_hdlr
|
|
from dsp.modules.hf import openai_to_hf
|
|
from dsp.modules.hf_client import send_hfvllm_request_v00, send_hftgi_request_v01_wrapped
|
|
from transformers import AutoTokenizer
|
|
|
|
try:
|
|
from anthropic import RateLimitError
|
|
except ImportError:
|
|
RateLimitError = None
|
|
|
|
|
|
class OpenAIModel(dspy.OpenAI):
|
|
"""A wrapper class for dspy.OpenAI."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str = "gpt-3.5-turbo-instruct",
|
|
api_key: Optional[str] = None,
|
|
model_type: Literal["chat", "text"] = None,
|
|
**kwargs
|
|
):
|
|
super().__init__(model=model, api_key=api_key, model_type=model_type, **kwargs)
|
|
self._token_usage_lock = threading.Lock()
|
|
self.prompt_tokens = 0
|
|
self.completion_tokens = 0
|
|
|
|
def log_usage(self, response):
|
|
"""Log the total tokens from the OpenAI API response."""
|
|
usage_data = response.get('usage')
|
|
if usage_data:
|
|
with self._token_usage_lock:
|
|
self.prompt_tokens += usage_data.get('prompt_tokens', 0)
|
|
self.completion_tokens += usage_data.get('completion_tokens', 0)
|
|
|
|
def get_usage_and_reset(self):
|
|
"""Get the total tokens used and reset the token usage."""
|
|
usage = {
|
|
self.kwargs.get('model') or self.kwargs.get('engine'):
|
|
{'prompt_tokens': self.prompt_tokens, 'completion_tokens': self.completion_tokens}
|
|
}
|
|
self.prompt_tokens = 0
|
|
self.completion_tokens = 0
|
|
|
|
return usage
|
|
|
|
def __call__(
|
|
self,
|
|
prompt: str,
|
|
only_completed: bool = True,
|
|
return_sorted: bool = False,
|
|
**kwargs,
|
|
) -> list[dict[str, Any]]:
|
|
"""Copied from dspy/dsp/modules/gpt3.py with the addition of tracking token usage."""
|
|
|
|
assert only_completed, "for now"
|
|
assert return_sorted is False, "for now"
|
|
|
|
# if kwargs.get("n", 1) > 1:
|
|
# if self.model_type == "chat":
|
|
# kwargs = {**kwargs}
|
|
# else:
|
|
# kwargs = {**kwargs, "logprobs": 5}
|
|
|
|
response = self.request(prompt, **kwargs)
|
|
|
|
# Log the token usage from the OpenAI API response.
|
|
self.log_usage(response)
|
|
|
|
choices = response["choices"]
|
|
|
|
completed_choices = [c for c in choices if c["finish_reason"] != "length"]
|
|
|
|
if only_completed and len(completed_choices):
|
|
choices = completed_choices
|
|
|
|
completions = [self._get_choice_text(c) for c in choices]
|
|
if return_sorted and kwargs.get("n", 1) > 1:
|
|
scored_completions = []
|
|
|
|
for c in choices:
|
|
tokens, logprobs = (
|
|
c["logprobs"]["tokens"],
|
|
c["logprobs"]["token_logprobs"],
|
|
)
|
|
|
|
if "<|endoftext|>" in tokens:
|
|
index = tokens.index("<|endoftext|>") + 1
|
|
tokens, logprobs = tokens[:index], logprobs[:index]
|
|
|
|
avglog = sum(logprobs) / len(logprobs)
|
|
scored_completions.append((avglog, self._get_choice_text(c)))
|
|
|
|
scored_completions = sorted(scored_completions, reverse=True)
|
|
completions = [c for _, c in scored_completions]
|
|
|
|
return completions
|
|
|
|
|
|
class AzureOpenAIModel(dspy.AzureOpenAI):
|
|
"""A wrapper class for dspy.AzureOpenAI."""
|
|
def __init__(
|
|
self,
|
|
api_base: Optional[str] = None,
|
|
api_version: Optional[str] = None,
|
|
model: str = "gpt-3.5-turbo-instruct",
|
|
api_key: Optional[str] = None,
|
|
model_type: Literal["chat", "text"] = "chat",
|
|
**kwargs,
|
|
):
|
|
super().__init__(
|
|
api_base=api_base, api_version=api_version, model=model, api_key=api_key, model_type=model_type, **kwargs)
|
|
self._token_usage_lock = threading.Lock()
|
|
self.prompt_tokens = 0
|
|
self.completion_tokens = 0
|
|
|
|
def log_usage(self, response):
|
|
"""Log the total tokens from the OpenAI API response.
|
|
Override log_usage() in dspy.AzureOpenAI for tracking accumulated token usage."""
|
|
usage_data = response.get('usage')
|
|
if usage_data:
|
|
with self._token_usage_lock:
|
|
self.prompt_tokens += usage_data.get('prompt_tokens', 0)
|
|
self.completion_tokens += usage_data.get('completion_tokens', 0)
|
|
|
|
def get_usage_and_reset(self):
|
|
"""Get the total tokens used and reset the token usage."""
|
|
usage = {
|
|
self.kwargs.get('model') or self.kwargs.get('engine'):
|
|
{'prompt_tokens': self.prompt_tokens, 'completion_tokens': self.completion_tokens}
|
|
}
|
|
self.prompt_tokens = 0
|
|
self.completion_tokens = 0
|
|
|
|
return usage
|
|
|
|
|
|
class ClaudeModel(dspy.dsp.modules.lm.LM):
|
|
"""Copied from dspy/dsp/modules/anthropic.py with the addition of tracking token usage."""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
api_key: Optional[str] = None,
|
|
api_base: Optional[str] = None,
|
|
**kwargs,
|
|
):
|
|
super().__init__(model)
|
|
try:
|
|
from anthropic import Anthropic
|
|
except ImportError as err:
|
|
raise ImportError("Claude requires `pip install anthropic`.") from err
|
|
|
|
self.provider = "anthropic"
|
|
self.api_key = api_key = os.environ.get("ANTHROPIC_API_KEY") if api_key is None else api_key
|
|
self.api_base = "https://api.anthropic.com/v1/messages" if api_base is None else api_base
|
|
self.kwargs = {"temperature": kwargs.get("temperature", 0.0),
|
|
"max_tokens": min(kwargs.get("max_tokens", 4096), 4096), "top_p": kwargs.get("top_p", 1.0),
|
|
"top_k": kwargs.get("top_k", 1), "n": kwargs.pop("n", kwargs.pop("num_generations", 1)),
|
|
**kwargs, "model": model}
|
|
self.history: list[dict[str, Any]] = []
|
|
self.client = Anthropic(api_key=api_key)
|
|
self.model = model
|
|
|
|
self._token_usage_lock = threading.Lock()
|
|
self.prompt_tokens = 0
|
|
self.completion_tokens = 0
|
|
|
|
def log_usage(self, response):
|
|
"""Log the total tokens from the Anthropic API response."""
|
|
usage_data = response.usage
|
|
if usage_data:
|
|
with self._token_usage_lock:
|
|
self.prompt_tokens += usage_data.input_tokens
|
|
self.completion_tokens += usage_data.output_tokens
|
|
|
|
def get_usage_and_reset(self):
|
|
"""Get the total tokens used and reset the token usage."""
|
|
usage = {
|
|
self.model:
|
|
{'prompt_tokens': self.prompt_tokens, 'completion_tokens': self.completion_tokens}
|
|
}
|
|
self.prompt_tokens = 0
|
|
self.completion_tokens = 0
|
|
|
|
return usage
|
|
|
|
def basic_request(self, prompt: str, **kwargs):
|
|
raw_kwargs = kwargs
|
|
kwargs = {**self.kwargs, **kwargs}
|
|
# caching mechanism requires hashable kwargs
|
|
kwargs["messages"] = [{"role": "user", "content": prompt}]
|
|
kwargs.pop("n")
|
|
response = self.client.messages.create(**kwargs)
|
|
# history = {
|
|
# "prompt": prompt,
|
|
# "response": response,
|
|
# "kwargs": kwargs,
|
|
# "raw_kwargs": raw_kwargs,
|
|
# }
|
|
json_serializable_history = {
|
|
"prompt": prompt,
|
|
"response": {
|
|
"content": response.content[0].text,
|
|
"model": response.model,
|
|
"role": response.role,
|
|
"stop_reason": response.stop_reason,
|
|
"stop_sequence": response.stop_sequence,
|
|
"type": response.type,
|
|
"usage": {
|
|
"input_tokens": response.usage.input_tokens,
|
|
"output_tokens": response.usage.output_tokens,
|
|
}
|
|
},
|
|
"kwargs": kwargs,
|
|
"raw_kwargs": raw_kwargs,
|
|
}
|
|
self.history.append(json_serializable_history)
|
|
return response
|
|
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
(RateLimitError,),
|
|
max_time=1000,
|
|
max_tries=8,
|
|
on_backoff=backoff_hdlr,
|
|
giveup=giveup_hdlr,
|
|
)
|
|
def request(self, prompt: str, **kwargs):
|
|
"""Handles retrieval of completions from Anthropic whilst handling API errors."""
|
|
return self.basic_request(prompt, **kwargs)
|
|
|
|
def __call__(self, prompt, only_completed=True, return_sorted=False, **kwargs):
|
|
"""Retrieves completions from Anthropic.
|
|
|
|
Args:
|
|
prompt (str): prompt to send to Anthropic
|
|
only_completed (bool, optional): return only completed responses and ignores completion due to length. Defaults to True.
|
|
return_sorted (bool, optional): sort the completion choices using the returned probabilities. Defaults to False.
|
|
|
|
Returns:
|
|
list[str]: list of completion choices
|
|
"""
|
|
assert only_completed, "for now"
|
|
assert return_sorted is False, "for now"
|
|
# per eg here: https://docs.anthropic.com/claude/reference/messages-examples
|
|
# max tokens can be used as a proxy to return smaller responses
|
|
# so this cannot be a proper indicator for incomplete response unless it isnt the user-intent.
|
|
n = kwargs.pop("n", 1)
|
|
completions = []
|
|
for _ in range(n):
|
|
response = self.request(prompt, **kwargs)
|
|
self.log_usage(response)
|
|
# This is the original behavior in dspy/dsp/modules/anthropic.py.
|
|
# Comment it out because it can cause "IndexError: list index out of range" silently
|
|
# which is not transparent to developers.
|
|
# if only_completed and response.stop_reason == "max_tokens":
|
|
# continue
|
|
completions = [c.text for c in response.content]
|
|
return completions
|
|
|
|
|
|
class VLLMClient(dspy.HFClientVLLM):
|
|
"""A wrapper class for dspy.HFClientVLLM."""
|
|
|
|
def __init__(self, model, port, url="http://localhost", **kwargs):
|
|
"""Copied from dspy/dsp/modules/hf_client.py with the addition of storing additional kwargs."""
|
|
|
|
super().__init__(model=model, port=port, url=url, **kwargs)
|
|
# Store additional kwargs for the generate method.
|
|
self.kwargs = {**self.kwargs, **kwargs}
|
|
|
|
def _generate(self, prompt, **kwargs):
|
|
"""Copied from dspy/dsp/modules/hf_client.py with the addition of passing kwargs to VLLM server."""
|
|
kwargs = {**self.kwargs, **kwargs}
|
|
|
|
# payload = {
|
|
# "model": kwargs["model"],
|
|
# "prompt": prompt,
|
|
# "max_tokens": kwargs["max_tokens"],
|
|
# "temperature": kwargs["temperature"],
|
|
# }
|
|
payload = {
|
|
"prompt": prompt,
|
|
**kwargs
|
|
}
|
|
|
|
response = send_hfvllm_request_v00(
|
|
f"{self.url}/v1/completions",
|
|
json=payload,
|
|
headers=self.headers,
|
|
)
|
|
|
|
try:
|
|
json_response = response.json()
|
|
completions = json_response["choices"]
|
|
response = {
|
|
"prompt": prompt,
|
|
"choices": [{"text": c["text"]} for c in completions],
|
|
}
|
|
return response
|
|
|
|
except Exception as e:
|
|
print("Failed to parse JSON response:", response.text)
|
|
raise Exception("Received invalid JSON response from server")
|
|
|
|
|
|
class OllamaClient(dspy.OllamaLocal):
|
|
"""A wrapper class for dspy.OllamaClient."""
|
|
|
|
def __init__(self, model, port, url="http://localhost", **kwargs):
|
|
"""Copied from dspy/dsp/modules/hf_client.py with the addition of storing additional kwargs."""
|
|
# Check if the URL has 'http://' or 'https://'
|
|
if not url.startswith("http://") and not url.startswith("https://"):
|
|
url = "http://" + url
|
|
super().__init__(model=model, base_url=f"{url}:{port}", **kwargs)
|
|
# Store additional kwargs for the generate method.
|
|
self.kwargs = {**self.kwargs, **kwargs}
|
|
|
|
|
|
class TGIClient(dspy.HFClientTGI):
|
|
def __init__(self, model, port, url, http_request_kwargs=None, **kwargs):
|
|
super().__init__(model=model, port=port, url=url, http_request_kwargs=http_request_kwargs, **kwargs)
|
|
|
|
def _generate(self, prompt, **kwargs):
|
|
"""Copied from dspy/dsp/modules/hf_client.py with the addition of removing hard-coded parameters."""
|
|
kwargs = {**self.kwargs, **kwargs}
|
|
|
|
payload = {
|
|
"inputs": prompt,
|
|
"parameters": {
|
|
"do_sample": kwargs["n"] > 1,
|
|
"best_of": kwargs["n"],
|
|
"details": kwargs["n"] > 1,
|
|
**kwargs,
|
|
},
|
|
}
|
|
|
|
payload["parameters"] = openai_to_hf(**payload["parameters"])
|
|
|
|
# Comment out the following lines to remove the hard-coded parameters.
|
|
# payload["parameters"]["temperature"] = max(
|
|
# 0.1, payload["parameters"]["temperature"],
|
|
# )
|
|
|
|
response = send_hftgi_request_v01_wrapped(
|
|
f"{self.url}:{random.Random().choice(self.ports)}" + "/generate",
|
|
url=self.url,
|
|
ports=tuple(self.ports),
|
|
json=payload,
|
|
headers=self.headers,
|
|
**self.http_request_kwargs,
|
|
)
|
|
|
|
try:
|
|
json_response = response.json()
|
|
# completions = json_response["generated_text"]
|
|
|
|
completions = [json_response["generated_text"]]
|
|
|
|
if (
|
|
"details" in json_response
|
|
and "best_of_sequences" in json_response["details"]
|
|
):
|
|
completions += [
|
|
x["generated_text"]
|
|
for x in json_response["details"]["best_of_sequences"]
|
|
]
|
|
|
|
response = {"prompt": prompt, "choices": [{"text": c} for c in completions]}
|
|
return response
|
|
except Exception:
|
|
print("Failed to parse JSON response:", response.text)
|
|
raise Exception("Received invalid JSON response from server")
|
|
|
|
|
|
class TogetherClient(dspy.HFModel):
|
|
"""A wrapper class for dspy.Together."""
|
|
|
|
def __init__(self, model, apply_tokenizer_chat_template=False, hf_tokenizer_name=None, **kwargs):
|
|
"""Copied from dspy/dsp/modules/hf_client.py with the support of applying tokenizer chat template."""
|
|
|
|
super().__init__(model=model, is_client=True)
|
|
self.session = requests.Session()
|
|
self.api_base = "https://api.together.xyz/v1/completions" if os.getenv(
|
|
"TOGETHER_API_BASE") is None else os.getenv("TOGETHER_API_BASE")
|
|
self.token = os.getenv("TOGETHER_API_KEY")
|
|
self.model = model
|
|
|
|
# self.use_inst_template = False
|
|
# if any(keyword in self.model.lower() for keyword in ["inst", "instruct"]):
|
|
# self.use_inst_template = True
|
|
self.apply_tokenizer_chat_template = apply_tokenizer_chat_template
|
|
if self.apply_tokenizer_chat_template:
|
|
logging.info("Loading huggingface tokenizer.")
|
|
if hf_tokenizer_name is None:
|
|
hf_tokenizer_name = self.model
|
|
self.tokenizer = AutoTokenizer.from_pretrained(hf_tokenizer_name, cache_dir=kwargs.get("cache_dir", None))
|
|
|
|
stop_default = "\n\n---"
|
|
|
|
self.kwargs = {
|
|
"temperature": 0.0,
|
|
"max_tokens": 512,
|
|
"top_p": 1,
|
|
"top_k": 20,
|
|
"repetition_penalty": 1,
|
|
"n": 1,
|
|
"stop": stop_default if "stop" not in kwargs else kwargs["stop"],
|
|
**kwargs,
|
|
}
|
|
self._token_usage_lock = threading.Lock()
|
|
self.prompt_tokens = 0
|
|
self.completion_tokens = 0
|
|
|
|
def log_usage(self, response):
|
|
"""Log the total tokens from the OpenAI API response."""
|
|
usage_data = response.get('usage')
|
|
if usage_data:
|
|
with self._token_usage_lock:
|
|
self.prompt_tokens += usage_data.get('prompt_tokens', 0)
|
|
self.completion_tokens += usage_data.get('completion_tokens', 0)
|
|
|
|
def get_usage_and_reset(self):
|
|
"""Get the total tokens used and reset the token usage."""
|
|
usage = {
|
|
self.model:
|
|
{'prompt_tokens': self.prompt_tokens, 'completion_tokens': self.completion_tokens}
|
|
}
|
|
self.prompt_tokens = 0
|
|
self.completion_tokens = 0
|
|
|
|
return usage
|
|
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
ERRORS,
|
|
max_time=1000,
|
|
on_backoff=backoff_hdlr,
|
|
)
|
|
def _generate(self, prompt, use_chat_api=False, **kwargs):
|
|
url = f"{self.api_base}"
|
|
|
|
kwargs = {**self.kwargs, **kwargs}
|
|
|
|
stop = kwargs.get("stop")
|
|
temperature = kwargs.get("temperature")
|
|
max_tokens = kwargs.get("max_tokens", 150)
|
|
top_p = kwargs.get("top_p", 0.7)
|
|
top_k = kwargs.get("top_k", 50)
|
|
repetition_penalty = kwargs.get("repetition_penalty", 1)
|
|
if self.apply_tokenizer_chat_template:
|
|
prompt = self.tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False)
|
|
# prompt = f"[INST]{prompt}[/INST]" if self.use_inst_template else prompt
|
|
|
|
if use_chat_api:
|
|
url = f"{self.api_base}/chat/completions"
|
|
messages = [
|
|
{"role": "system",
|
|
"content": "You are a helpful assistant. You must continue the user text directly without *any* additional interjections."},
|
|
{"role": "user", "content": prompt},
|
|
]
|
|
body = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"repetition_penalty": repetition_penalty,
|
|
"stop": stop,
|
|
}
|
|
else:
|
|
body = {
|
|
"model": self.model,
|
|
"prompt": prompt,
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"repetition_penalty": repetition_penalty,
|
|
"stop": stop,
|
|
}
|
|
|
|
headers = {"Authorization": f"Bearer {self.token}"}
|
|
|
|
with self.session.post(url, headers=headers, json=body) as resp:
|
|
resp_json = resp.json()
|
|
# Log the token usage from the Together API response.
|
|
self.log_usage(resp_json)
|
|
if use_chat_api:
|
|
# completions = [resp_json['output'].get('choices', [])[0].get('message', {}).get('content', "")]
|
|
completions = [resp_json.get('choices', [])[0].get('message', {}).get('content', "")]
|
|
else:
|
|
# completions = [resp_json['output'].get('choices', [])[0].get('text', "")]
|
|
completions = [resp_json.get('choices', [])[0].get('text', "")]
|
|
response = {"prompt": prompt, "choices": [{"text": c} for c in completions]}
|
|
return response
|