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: