scoring function def rename

This commit is contained in:
Xi Yan 2024-10-24 13:51:11 -07:00
parent 42bac85e1f
commit 6053b8dd34
4 changed files with 19 additions and 40 deletions

View file

@ -4,20 +4,10 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
from typing import ( from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
Any,
Dict,
List,
Literal,
Optional,
Protocol,
runtime_checkable,
Union,
)
from llama_models.schema_utils import json_schema_type, webmethod from llama_models.schema_utils import json_schema_type, webmethod
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from typing_extensions import Annotated
from llama_stack.apis.common.type_system import ParamType from llama_stack.apis.common.type_system import ParamType
@ -33,21 +23,19 @@ class Parameter(BaseModel):
# with standard metrics so they can be rolled up? # with standard metrics so they can be rolled up?
class LLMAsJudgeContext(BaseModel):
judge_model: str
prompt_template: Optional[str] = None
@json_schema_type @json_schema_type
class CommonFunctionDef(BaseModel): class ScoringFunctionDef(BaseModel):
identifier: str identifier: str
description: Optional[str] = None description: Optional[str] = None
metadata: Dict[str, Any] = Field( metadata: Dict[str, Any] = Field(
default_factory=dict, default_factory=dict,
description="Any additional metadata for this definition", description="Any additional metadata for this definition",
) )
# Hack: same with memory_banks for union defs
provider_id: str = ""
@json_schema_type
class DeterministicFunctionDef(CommonFunctionDef):
type: Literal["deterministic"] = "deterministic"
parameters: List[Parameter] = Field( parameters: List[Parameter] = Field(
description="List of parameters for the deterministic function", description="List of parameters for the deterministic function",
default_factory=list, default_factory=list,
@ -55,24 +43,17 @@ class DeterministicFunctionDef(CommonFunctionDef):
return_type: ParamType = Field( return_type: ParamType = Field(
description="The return type of the deterministic function", description="The return type of the deterministic function",
) )
context: Optional[LLMAsJudgeContext] = None
# We can optionally add information here to support packaging of code, etc. # We can optionally add information here to support packaging of code, etc.
@json_schema_type @json_schema_type
class LLMJudgeFunctionDef(CommonFunctionDef): class ScoringFunctionDefWithProvider(ScoringFunctionDef):
type: Literal["judge"] = "judge" provider_id: str = Field(
model: str = Field( description="ID of the provider which serves this dataset",
description="The LLM model to use for the judge function",
) )
ScoringFunctionDef = Annotated[
Union[DeterministicFunctionDef, LLMJudgeFunctionDef], Field(discriminator="type")
]
ScoringFunctionDefWithProvider = ScoringFunctionDef
@runtime_checkable @runtime_checkable
class ScoringFunctions(Protocol): class ScoringFunctions(Protocol):
@webmethod(route="/scoring_functions/list", method="GET") @webmethod(route="/scoring_functions/list", method="GET")

View file

@ -95,17 +95,15 @@ class CommonRoutingTableImpl(RoutingTable):
for d in datasets: for d in datasets:
d.provider_id = pid d.provider_id = pid
add_objects(datasets)
elif api == Api.scoring: elif api == Api.scoring:
p.scoring_function_store = self p.scoring_function_store = self
scoring_functions = await p.list_scoring_functions() scoring_functions = await p.list_scoring_functions()
add_objects(
# do in-memory updates due to pesky Annotated unions [
for s in scoring_functions: ScoringFunctionDefWithProvider(**s.dict(), provider_id=pid)
s.provider_id = pid for s in scoring_functions
]
add_objects(scoring_functions) )
async def shutdown(self) -> None: async def shutdown(self) -> None:
for p in self.impls_by_provider_id.values(): for p in self.impls_by_provider_id.values():

View file

@ -17,7 +17,7 @@ class BaseScorer(ABC):
- aggregate(self, scorer_results) - aggregate(self, scorer_results)
""" """
scoring_function_def: DeterministicFunctionDef scoring_function_def: ScoringFunctionDef
def __init__(self, *args, **kwargs) -> None: def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)

View file

@ -17,7 +17,7 @@ class EqualityScorer(BaseScorer):
A scorer that assigns a score of 1.0 if the input string matches the target string, and 0.0 otherwise. A scorer that assigns a score of 1.0 if the input string matches the target string, and 0.0 otherwise.
""" """
scoring_function_def = DeterministicFunctionDef( scoring_function_def = ScoringFunctionDef(
identifier="equality", identifier="equality",
description="Returns 1.0 if the input is equal to the target, 0.0 otherwise.", description="Returns 1.0 if the input is equal to the target, 0.0 otherwise.",
parameters=[], parameters=[],