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
131 lines
3.1 KiB
Python
131 lines
3.1 KiB
Python
import abc
|
|
import enum
|
|
from typing import Any, TypeVar
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from tau_bench.model_utils.api.datapoint import (
|
|
BinaryClassifyDatapoint,
|
|
ClassifyDatapoint,
|
|
Datapoint,
|
|
GenerateDatapoint,
|
|
ParseDatapoint,
|
|
ParseForceDatapoint,
|
|
ScoreDatapoint,
|
|
)
|
|
from tau_bench.model_utils.api.types import PartialObj
|
|
|
|
T = TypeVar("T", bound=BaseModel)
|
|
|
|
|
|
class Platform(enum.Enum):
|
|
OPENAI = "openai"
|
|
MISTRAL = "mistral"
|
|
ANTHROPIC = "anthropic"
|
|
ANYSCALE = "anyscale"
|
|
OUTLINES = "outlines"
|
|
VLLM_CHAT = "vllm-chat"
|
|
VLLM_COMPLETION = "vllm-completion"
|
|
|
|
|
|
# @runtime_checkable
|
|
# class Model(Protocol):
|
|
class Model(abc.ABC):
|
|
@abc.abstractmethod
|
|
def get_capability(self) -> float:
|
|
"""Return the capability of the model, a float between 0.0 and 1.0."""
|
|
raise NotImplementedError
|
|
|
|
@abc.abstractmethod
|
|
def get_approx_cost(self, dp: Datapoint) -> float:
|
|
raise NotImplementedError
|
|
|
|
@abc.abstractmethod
|
|
def get_latency(self, dp: Datapoint) -> float:
|
|
raise NotImplementedError
|
|
|
|
@abc.abstractmethod
|
|
def supports_dp(self, dp: Datapoint) -> bool:
|
|
raise NotImplementedError
|
|
|
|
|
|
class ClassifyModel(Model):
|
|
@abc.abstractmethod
|
|
def classify(
|
|
self,
|
|
instruction: str,
|
|
text: str,
|
|
options: list[str],
|
|
examples: list[ClassifyDatapoint] | None = None,
|
|
temperature: float | None = None,
|
|
) -> int:
|
|
raise NotImplementedError
|
|
|
|
|
|
class BinaryClassifyModel(Model):
|
|
@abc.abstractmethod
|
|
def binary_classify(
|
|
self,
|
|
instruction: str,
|
|
text: str,
|
|
examples: list[BinaryClassifyDatapoint] | None = None,
|
|
temperature: float | None = None,
|
|
) -> bool:
|
|
raise NotImplementedError
|
|
|
|
|
|
class ParseModel(Model):
|
|
@abc.abstractmethod
|
|
def parse(
|
|
self,
|
|
text: str,
|
|
typ: type[T] | dict[str, Any],
|
|
examples: list[ParseDatapoint] | None = None,
|
|
temperature: float | None = None,
|
|
) -> T | PartialObj | dict[str, Any]:
|
|
raise NotImplementedError
|
|
|
|
|
|
class GenerateModel(Model):
|
|
@abc.abstractmethod
|
|
def generate(
|
|
self,
|
|
instruction: str,
|
|
text: str,
|
|
examples: list[GenerateDatapoint] | None = None,
|
|
temperature: float | None = None,
|
|
) -> str:
|
|
raise NotImplementedError
|
|
|
|
|
|
class ParseForceModel(Model):
|
|
@abc.abstractmethod
|
|
def parse_force(
|
|
self,
|
|
instruction: str,
|
|
typ: type[T] | dict[str, Any],
|
|
text: str | None = None,
|
|
examples: list[ParseForceDatapoint] | None = None,
|
|
temperature: float | None = None,
|
|
) -> T | dict[str, Any]:
|
|
raise NotImplementedError
|
|
|
|
|
|
class ScoreModel(Model):
|
|
@abc.abstractmethod
|
|
def score(
|
|
self,
|
|
instruction: str,
|
|
text: str,
|
|
min: int,
|
|
max: int,
|
|
examples: list[ScoreDatapoint] | None = None,
|
|
temperature: float | None = None,
|
|
) -> int:
|
|
raise NotImplementedError
|
|
|
|
|
|
AnyModel = (
|
|
BinaryClassifyModel | ClassifyModel | ParseForceModel | GenerateModel | ParseModel | ScoreModel
|
|
)
|