diff --git a/source/api_definitions.py b/source/api_definitions.py index a1ea720c8..a02536453 100644 --- a/source/api_definitions.py +++ b/source/api_definitions.py @@ -29,6 +29,7 @@ from model_types import ( InstructModel, Message, PretrainedModel, + RewardModel, SamplingParams, ShieldConfig, StopReason, @@ -245,9 +246,7 @@ class RewardScoringRequest: """Request to score a reward function. A list of prompts and a list of responses per prompt.""" prompt_generations: List[KPromptGenerations] - - # TODO(ragho): create a RewardModel enum tye - model: str + model: RewardModel @json_schema_type diff --git a/source/model_types.py b/source/model_types.py index 8ee6cef3b..f695938f3 100644 --- a/source/model_types.py +++ b/source/model_types.py @@ -137,3 +137,6 @@ class PretrainedModel(Enum): class InstructModel(Enum): llama3_8b_chat = "llama3_8b_chat" llama3_70b_chat = "llama3_70b_chat" + +class RewardModel(Enum): + llama3_405b_reward = "llama3_405b_reward" diff --git a/source/openapi.html b/source/openapi.html index 63994dd77..6e51f862a 100644 --- a/source/openapi.html +++ b/source/openapi.html @@ -1976,7 +1976,10 @@ } }, "model": { - "type": "string" + "type": "string", + "enum": [ + "llama3_405b_reward" + ] } }, "additionalProperties": false, @@ -2312,23 +2315,23 @@ } ], "tags": [ + { + "name": "RewardScoring" + }, { "name": "Inference" }, - { - "name": "SyntheticDataGeneration" - }, - { - "name": "Datasets" - }, { "name": "AgenticSystem" }, { - "name": "Finetuning" + "name": "Datasets" }, { - "name": "RewardScoring" + "name": "SyntheticDataGeneration" + }, + { + "name": "Finetuning" }, { "name": "ShieldConfig", diff --git a/source/openapi.yaml b/source/openapi.yaml index 472b37196..9acc76086 100644 --- a/source/openapi.yaml +++ b/source/openapi.yaml @@ -1039,6 +1039,8 @@ components: additionalProperties: false properties: model: + enum: + - llama3_405b_reward type: string prompt_generations: items: @@ -1412,12 +1414,12 @@ security: servers: - url: http://llama.meta.com tags: -- name: Inference -- name: SyntheticDataGeneration -- name: Datasets -- name: AgenticSystem -- name: Finetuning - name: RewardScoring +- name: Inference +- name: AgenticSystem +- name: Datasets +- name: SyntheticDataGeneration +- name: Finetuning - description: name: ShieldConfig - description: