liqiang
b119135836
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
122 lines
4.4 KiB
Python
122 lines
4.4 KiB
Python
import os
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from tau_bench.model_utils.api.datapoint import Datapoint
|
|
from tau_bench.model_utils.model.completion import (
|
|
CompletionModel,
|
|
approx_cost_for_datapoint,
|
|
approx_prompt_str,
|
|
)
|
|
from tau_bench.model_utils.model.utils import approx_num_tokens
|
|
from tau_bench.model_utils.model.vllm_utils import generate_request
|
|
|
|
PRICE_PER_INPUT_TOKEN_MAP = {
|
|
"Qwen/Qwen2-0.5B-Instruct": 0.0,
|
|
"Qwen/Qwen2-1.5B-Instruct": 0.0,
|
|
"Qwen/Qwen2-7B-Instruct": 0.0,
|
|
"Qwen/Qwen2-72B-Instruct": 0.0,
|
|
"meta-llama/Meta-Llama-3-8B-Instruct": 0.0,
|
|
"meta-llama/Meta-Llama-3.1-8B-Instruct": 0.0,
|
|
"meta-llama/Meta-Llama-3-70B-Instruct": 0.0,
|
|
"mistralai/Mistral-Nemo-Instruct-2407": 0.0,
|
|
}
|
|
INPUT_PRICE_PER_TOKEN_FALLBACK = 0.0
|
|
|
|
# TODO: refine this
|
|
CAPABILITY_SCORE_MAP = {
|
|
"Qwen/Qwen2-0.5B-Instruct": 0.05,
|
|
"Qwen/Qwen2-1.5B-Instruct": 0.07,
|
|
"Qwen/Qwen2-7B-Instruct": 0.2,
|
|
"Qwen/Qwen2-72B-Instruct": 0.4,
|
|
"meta-llama/Meta-Llama-3.1-8B-Instruct": 0.3,
|
|
"sierra-research/Meta-Llama-3.1-8B-Instruct": 0.3,
|
|
"meta-llama/Meta-Llama-3.1-70B-Instruct": 0.5,
|
|
"mistralai/Mistral-Nemo-Instruct-2407": 0.3,
|
|
}
|
|
CAPABILITY_SCORE_FALLBACK = 0.1
|
|
|
|
# TODO: implement
|
|
LATENCY_MS_PER_OUTPUT_TOKEN_MAP = {}
|
|
# TODO: implement
|
|
LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK = 0.0
|
|
|
|
MAX_CONTEXT_LENGTH_MAP = {
|
|
"Qwen/Qwen2-0.5B-Instruct": 32768,
|
|
"Qwen/Qwen2-1.5B-Instruct": 32768,
|
|
"Qwen/Qwen2-7B-Instruct": 131072,
|
|
"Qwen/Qwen2-72B-Instruct": 131072,
|
|
"meta-llama/Meta-Llama-3.1-8B-Instruct": 128000,
|
|
"sierra-research/Meta-Llama-3.1-8B-Instruct": 128000,
|
|
"meta-llama/Meta-Llama-3.1-70B-Instruct": 128000,
|
|
"mistralai/Mistral-Nemo-Instruct-2407": 128000,
|
|
}
|
|
MAX_CONTEXT_LENGTH_FALLBACK = 128000
|
|
|
|
|
|
class VLLMCompletionModel(CompletionModel):
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
base_url: str,
|
|
endpoint: str = "generate",
|
|
temperature: float = 0.0,
|
|
price_per_input_token: float | None = None,
|
|
capability: float | None = None,
|
|
latency_ms_per_output_token: float | None = None,
|
|
max_context_length: int | None = None,
|
|
) -> None:
|
|
self.model = model
|
|
self.base_url = base_url
|
|
self.url = os.path.join(base_url, endpoint)
|
|
self.temperature = temperature
|
|
self.price_per_input_token = (
|
|
price_per_input_token
|
|
if price_per_input_token is not None
|
|
else PRICE_PER_INPUT_TOKEN_MAP.get(model, INPUT_PRICE_PER_TOKEN_FALLBACK)
|
|
)
|
|
self.capability = (
|
|
capability
|
|
if capability is not None
|
|
else CAPABILITY_SCORE_MAP.get(model, CAPABILITY_SCORE_FALLBACK)
|
|
)
|
|
self.latency_ms_per_output_token = (
|
|
latency_ms_per_output_token
|
|
if latency_ms_per_output_token is not None
|
|
else LATENCY_MS_PER_OUTPUT_TOKEN_MAP.get(model, LATENCY_MS_PER_OUTPUT_TOKEN_FALLBACK)
|
|
)
|
|
self.max_context_length = (
|
|
max_context_length
|
|
if max_context_length is not None
|
|
else MAX_CONTEXT_LENGTH_MAP.get(model, MAX_CONTEXT_LENGTH_FALLBACK)
|
|
)
|
|
|
|
def generate_from_prompt(self, prompt: str, temperature: float = 0.0) -> str:
|
|
return generate_request(url=self.url, prompt=prompt, temperature=temperature)
|
|
|
|
def parse_force_from_prompt(
|
|
self, prompt: str, typ: BaseModel | dict[str, Any], temperature: float | None = None
|
|
) -> dict[str, Any]:
|
|
if temperature is None:
|
|
temperature = self.temperature
|
|
res = generate_request(
|
|
url=self.url, prompt=prompt, force_json=True, temperature=temperature
|
|
)
|
|
return self.handle_parse_force_response(prompt=prompt, content=res)
|
|
|
|
def get_approx_cost(self, dp: Datapoint) -> float:
|
|
cost_per_token = self.price_per_input_token
|
|
return approx_cost_for_datapoint(dp=dp, price_per_input_token=cost_per_token)
|
|
|
|
def get_latency(self, dp: Datapoint) -> float:
|
|
latency_per_output_token = self.latency_ms_per_output_token
|
|
return approx_cost_for_datapoint(dp=dp, price_per_input_token=latency_per_output_token)
|
|
|
|
def get_capability(self) -> float:
|
|
return CAPABILITY_SCORE_MAP.get(self.model, CAPABILITY_SCORE_FALLBACK)
|
|
|
|
def supports_dp(self, dp: Datapoint) -> bool:
|
|
prompt = approx_prompt_str(dp)
|
|
return approx_num_tokens(prompt) <= self.max_context_length
|