diff --git a/source/api_definitions.py b/source/api_definitions.py
index 4dded2166..d8c144415 100644
--- a/source/api_definitions.py
+++ b/source/api_definitions.py
@@ -1,6 +1,7 @@
from dataclasses import dataclass, field
+from datetime import datetime
from enum import Enum
-from typing import Any, Dict, List, Optional, Protocol, Set, Union, Tuple
+from typing import Any, Dict, List, Optional, Protocol, Set, Tuple, Union
import yaml
from agentic_system_types import (
@@ -10,6 +11,18 @@ from agentic_system_types import (
SafetyViolation,
)
+from finetuning_types import (
+ Dataset,
+ DoraFinetuningConfig,
+ FinetuningAlgorithm,
+ FinetuningJobLogStream,
+ FinetuningJobStatus,
+ LoraFinetuningConfig,
+ OptimizerConfig,
+ QLoraFinetuningConfig,
+ TrainingConfig,
+)
+
from model_types import (
BuiltinTool,
Content,
@@ -22,6 +35,7 @@ from model_types import (
ToolCall,
ToolDefinition,
ToolResponse,
+ URL,
)
from pyopenapi import Info, Options, Server, Specification, webmethod
@@ -205,6 +219,7 @@ class AgenticSystem(Protocol):
@dataclass
class PromptGeneration:
+ # TODO(ashwin): probably create a Dialog type which is used everywhere including chat completion
prompt: Message
message_history: List[Message]
generation: Message
@@ -286,8 +301,99 @@ class SyntheticDataGeneration(Protocol):
) -> Union[SyntheticDataGenerationResponse]: ...
+@json_schema_type
+@dataclass
+class CreateDatasetRequest:
+ """Request to create a dataset."""
+
+ uuid: str
+ dataset: Dataset
+
+
+class Datasets(Protocol):
+ @webmethod(route="/datasets/create")
+ def create_dataset(
+ self,
+ request: CreateDatasetRequest,
+ ) -> None: ...
+
+ @webmethod(route="/datasets/get")
+ def get_dataset(
+ self,
+ dataset_id: str,
+ ) -> Dataset: ...
+
+ @webmethod(route="/datasets/delete")
+ def delete_dataset(
+ self,
+ dataset_id: str,
+ ) -> None: ...
+
+
+@json_schema_type
+@dataclass
+class FinetuningTrainRequest:
+ """Request to finetune a model."""
+
+ job_uuid: str
+
+ model: PretrainedModel
+ dataset: Dataset
+ validation_dataset: Dataset
+
+ algorithm: FinetuningAlgorithm
+ algorithm_config: Union[
+ LoraFinetuningConfig, QLoraFinetuningConfig, DoraFinetuningConfig
+ ]
+
+ optimizer_config: OptimizerConfig
+ training_config: TrainingConfig
+
+ # TODO: define these
+ hyperparam_search_config: Dict[str, Any]
+ logger_config: Dict[str, Any]
+
+
+@json_schema_type
+@dataclass
+class FinetuningJobStatusResponse:
+ """Status of a finetuning job."""
+
+ job_uuid: str
+ status: FinetuningJobStatus
+
+ scheduled_at: Optional[datetime] = None
+ started_at: Optional[datetime] = None
+ completed_at: Optional[datetime] = None
+
+ resources_allocated: Optional[Dict[str, Any]] = None
+
+
+class Finetuning(Protocol):
+ @webmethod(route="/finetuning/text_generation/train")
+ def post_train(
+ self,
+ request: FinetuningTrainRequest,
+ ) -> None: ...
+
+ # sends SSE stream of logs
+ @webmethod(route="/finetuning/job/logs")
+ def get_training_log_stream(self, job_uuid: str) -> FinetuningJobLogStream: ...
+
+ @webmethod(route="/finetuning/job/status")
+ def get_training_job_status(self, job_uuid: str) -> FinetuningJobStatusResponse: ...
+
+ @webmethod(route="/finetuning/job/cancel")
+ def cancel_training_job(self, job_uuid: str) -> None: ...
+
+
class LlamaStackEndpoints(
- Inference, AgenticSystem, RewardScoring, SyntheticDataGeneration
+ Inference,
+ AgenticSystem,
+ RewardScoring,
+ SyntheticDataGeneration,
+ Datasets,
+ Finetuning,
): ...
diff --git a/source/finetuning_types.py b/source/finetuning_types.py
new file mode 100644
index 000000000..99c3f112d
--- /dev/null
+++ b/source/finetuning_types.py
@@ -0,0 +1,98 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional, Set, Union
+
+from model_types import Message, URL
+
+from strong_typing.schema import json_schema_type
+
+
+class DatasetColumnType(Enum):
+ dialog = "dialog"
+ text = "text"
+ media = "media"
+ number = "number"
+ json = "json"
+
+
+@json_schema_type
+@dataclass
+class Dataset:
+ """Dataset to be used for training or evaluating language models."""
+
+ # TODO(ashwin): figure out if we need to add an enum for a "dataset type"
+
+ columns: Dict[str, DatasetColumnType]
+ content_url: URL
+ metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+class OptimizerType(Enum):
+ adam = "adam"
+ adamw = "adamw"
+ sgd = "sgd"
+
+
+@json_schema_type
+@dataclass
+class OptimizerConfig:
+ optimizer_type: OptimizerType
+ lr: float
+ lr_min: float
+ weight_decay: float
+
+
+@json_schema_type
+@dataclass
+class TrainingConfig:
+ n_epochs: int
+ batch_size: int
+ shuffle: bool
+ n_iters: int
+
+ enable_activation_checkpointing: bool
+ memory_efficient_fsdp_wrap: bool
+ fsdp_cpu_offload: bool
+
+
+class FinetuningAlgorithm(Enum):
+ full = "full"
+ lora = "lora"
+ qlora = "qlora"
+ dora = "dora"
+
+
+@json_schema_type
+@dataclass
+class LoraFinetuningConfig:
+ lora_attn_modules: List[str]
+ apply_lora_to_mlp: bool
+ apply_lora_to_output: bool
+ rank: int
+ alpha: int
+
+
+@dataclass
+class QLoraFinetuningConfig(LoraFinetuningConfig):
+ pass
+
+
+@dataclass
+class DoraFinetuningConfig(LoraFinetuningConfig):
+ pass
+
+
+@json_schema_type
+@dataclass
+class FinetuningJobLogStream:
+ """Stream of logs from a finetuning job."""
+
+ job_uuid: str
+ log_lines: List[str]
+
+
+class FinetuningJobStatus(Enum):
+ running = "running"
+ completed = "completed"
+ failed = "failed"
+ scheduled = "scheduled"
diff --git a/source/model_types.py b/source/model_types.py
index 795c0b462..8ee6cef3b 100644
--- a/source/model_types.py
+++ b/source/model_types.py
@@ -49,7 +49,7 @@ class Attachment:
url: URL
mime_type: str
-
+# TODO(ashwin): make this better named maybe InterleavedTextMedia
Content = Union[
str,
Attachment,
@@ -77,7 +77,7 @@ class ToolCall:
@dataclass
class ToolResponse:
tool_name: str
- response: str
+ content: Content
# TODO: we need to document the parameters for the tool calls
diff --git a/source/openapi.html b/source/openapi.html
index 8f5f935d3..987897354 100644
--- a/source/openapi.html
+++ b/source/openapi.html
@@ -96,6 +96,29 @@
}
}
},
+ "/datasets/create": {
+ "post": {
+ "responses": {
+ "200": {
+ "description": "OK"
+ }
+ },
+ "tags": [
+ "Datasets"
+ ],
+ "parameters": [],
+ "requestBody": {
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/CreateDatasetRequest"
+ }
+ }
+ },
+ "required": true
+ }
+ }
+ },
"/agentic_system/delete": {
"delete": {
"responses": {
@@ -118,6 +141,115 @@
]
}
},
+ "/datasets/delete": {
+ "delete": {
+ "responses": {
+ "200": {
+ "description": "OK"
+ }
+ },
+ "tags": [
+ "Datasets"
+ ],
+ "parameters": [
+ {
+ "name": "dataset_id",
+ "in": "query",
+ "required": true,
+ "schema": {
+ "type": "string"
+ }
+ }
+ ]
+ }
+ },
+ "/datasets/get": {
+ "get": {
+ "responses": {
+ "200": {
+ "description": "OK",
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/Dataset"
+ }
+ }
+ }
+ }
+ },
+ "tags": [
+ "Datasets"
+ ],
+ "parameters": [
+ {
+ "name": "dataset_id",
+ "in": "query",
+ "required": true,
+ "schema": {
+ "type": "string"
+ }
+ }
+ ]
+ }
+ },
+ "/finetuning/job/status": {
+ "get": {
+ "responses": {
+ "200": {
+ "description": "OK",
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/FinetuningJobStatusResponse"
+ }
+ }
+ }
+ }
+ },
+ "tags": [
+ "Finetuning"
+ ],
+ "parameters": [
+ {
+ "name": "job_uuid",
+ "in": "query",
+ "required": true,
+ "schema": {
+ "type": "string"
+ }
+ }
+ ]
+ }
+ },
+ "/finetuning/job/logs": {
+ "get": {
+ "responses": {
+ "200": {
+ "description": "OK",
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/FinetuningJobLogStream"
+ }
+ }
+ }
+ }
+ },
+ "tags": [
+ "Finetuning"
+ ],
+ "parameters": [
+ {
+ "name": "job_uuid",
+ "in": "query",
+ "required": true,
+ "schema": {
+ "type": "string"
+ }
+ }
+ ]
+ }
+ },
"/chat_completion": {
"post": {
"responses": {
@@ -251,6 +383,29 @@
"required": true
}
}
+ },
+ "/finetuning/text_generation/train": {
+ "post": {
+ "responses": {
+ "200": {
+ "description": "OK"
+ }
+ },
+ "tags": [
+ "Finetuning"
+ ],
+ "parameters": [],
+ "requestBody": {
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/FinetuningTrainRequest"
+ }
+ }
+ },
+ "required": true
+ }
+ }
}
},
"jsonSchemaDialect": "https://json-schema.org/draft/2020-12/schema",
@@ -586,14 +741,34 @@
"tool_name": {
"type": "string"
},
- "response": {
- "type": "string"
+ "content": {
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "$ref": "#/components/schemas/Attachment"
+ },
+ {
+ "type": "array",
+ "items": {
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "$ref": "#/components/schemas/Attachment"
+ }
+ ]
+ }
+ }
+ ]
}
},
"additionalProperties": false,
"required": [
"tool_name",
- "response"
+ "content"
]
}
}
@@ -817,14 +992,34 @@
"tool_name": {
"type": "string"
},
- "response": {
- "type": "string"
+ "content": {
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "$ref": "#/components/schemas/Attachment"
+ },
+ {
+ "type": "array",
+ "items": {
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "$ref": "#/components/schemas/Attachment"
+ }
+ ]
+ }
+ }
+ ]
}
},
"additionalProperties": false,
"required": [
"tool_name",
- "response"
+ "content"
]
}
}
@@ -944,14 +1139,34 @@
"tool_name": {
"type": "string"
},
- "response": {
- "type": "string"
+ "content": {
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "$ref": "#/components/schemas/Attachment"
+ },
+ {
+ "type": "array",
+ "items": {
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "$ref": "#/components/schemas/Attachment"
+ }
+ ]
+ }
+ }
+ ]
}
},
"additionalProperties": false,
"required": [
"tool_name",
- "response"
+ "content"
]
},
"response_text_delta": {
@@ -991,6 +1206,156 @@
],
"title": "Streamed agent execution response."
},
+ "CreateDatasetRequest": {
+ "type": "object",
+ "properties": {
+ "uuid": {
+ "type": "string"
+ },
+ "dataset": {
+ "$ref": "#/components/schemas/Dataset"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "uuid",
+ "dataset"
+ ],
+ "title": "Request to create a dataset."
+ },
+ "Dataset": {
+ "type": "object",
+ "properties": {
+ "columns": {
+ "type": "object",
+ "additionalProperties": {
+ "type": "string",
+ "enum": [
+ "dialog",
+ "text",
+ "media",
+ "number",
+ "json"
+ ]
+ }
+ },
+ "content_url": {
+ "$ref": "#/components/schemas/URL"
+ },
+ "metadata": {
+ "type": "object",
+ "additionalProperties": {
+ "oneOf": [
+ {
+ "type": "null"
+ },
+ {
+ "type": "boolean"
+ },
+ {
+ "type": "number"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "array"
+ },
+ {
+ "type": "object"
+ }
+ ]
+ }
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "columns",
+ "content_url",
+ "metadata"
+ ],
+ "title": "Dataset to be used for training or evaluating language models."
+ },
+ "FinetuningJobStatusResponse": {
+ "type": "object",
+ "properties": {
+ "job_uuid": {
+ "type": "string"
+ },
+ "status": {
+ "type": "string",
+ "enum": [
+ "running",
+ "completed",
+ "failed",
+ "scheduled"
+ ]
+ },
+ "scheduled_at": {
+ "type": "string",
+ "format": "date-time"
+ },
+ "started_at": {
+ "type": "string",
+ "format": "date-time"
+ },
+ "completed_at": {
+ "type": "string",
+ "format": "date-time"
+ },
+ "resources_allocated": {
+ "type": "object",
+ "additionalProperties": {
+ "oneOf": [
+ {
+ "type": "null"
+ },
+ {
+ "type": "boolean"
+ },
+ {
+ "type": "number"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "array"
+ },
+ {
+ "type": "object"
+ }
+ ]
+ }
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "job_uuid",
+ "status"
+ ],
+ "title": "Status of a finetuning job."
+ },
+ "FinetuningJobLogStream": {
+ "type": "object",
+ "properties": {
+ "job_uuid": {
+ "type": "string"
+ },
+ "log_lines": {
+ "type": "array",
+ "items": {
+ "type": "string"
+ }
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "job_uuid",
+ "log_lines"
+ ],
+ "title": "Stream of logs from a finetuning job."
+ },
"ChatCompletionRequest": {
"type": "object",
"properties": {
@@ -1666,6 +2031,271 @@
"scored_generations"
],
"title": "Response from the reward scoring. Batch of (prompt, response, score) tuples that pass the threshold."
+ },
+ "FinetuningTrainRequest": {
+ "type": "object",
+ "properties": {
+ "job_uuid": {
+ "type": "string"
+ },
+ "model": {
+ "type": "string",
+ "enum": [
+ "llama3_8b",
+ "llama3_70b"
+ ]
+ },
+ "dataset": {
+ "$ref": "#/components/schemas/Dataset"
+ },
+ "validation_dataset": {
+ "$ref": "#/components/schemas/Dataset"
+ },
+ "algorithm": {
+ "type": "string",
+ "enum": [
+ "full",
+ "lora",
+ "qlora",
+ "dora"
+ ]
+ },
+ "algorithm_config": {
+ "oneOf": [
+ {
+ "$ref": "#/components/schemas/LoraFinetuningConfig"
+ },
+ {
+ "type": "object",
+ "properties": {
+ "lora_attn_modules": {
+ "type": "array",
+ "items": {
+ "type": "string"
+ }
+ },
+ "apply_lora_to_mlp": {
+ "type": "boolean"
+ },
+ "apply_lora_to_output": {
+ "type": "boolean"
+ },
+ "rank": {
+ "type": "integer"
+ },
+ "alpha": {
+ "type": "integer"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "lora_attn_modules",
+ "apply_lora_to_mlp",
+ "apply_lora_to_output",
+ "rank",
+ "alpha"
+ ]
+ },
+ {
+ "type": "object",
+ "properties": {
+ "lora_attn_modules": {
+ "type": "array",
+ "items": {
+ "type": "string"
+ }
+ },
+ "apply_lora_to_mlp": {
+ "type": "boolean"
+ },
+ "apply_lora_to_output": {
+ "type": "boolean"
+ },
+ "rank": {
+ "type": "integer"
+ },
+ "alpha": {
+ "type": "integer"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "lora_attn_modules",
+ "apply_lora_to_mlp",
+ "apply_lora_to_output",
+ "rank",
+ "alpha"
+ ]
+ }
+ ]
+ },
+ "optimizer_config": {
+ "$ref": "#/components/schemas/OptimizerConfig"
+ },
+ "training_config": {
+ "$ref": "#/components/schemas/TrainingConfig"
+ },
+ "hyperparam_search_config": {
+ "type": "object",
+ "additionalProperties": {
+ "oneOf": [
+ {
+ "type": "null"
+ },
+ {
+ "type": "boolean"
+ },
+ {
+ "type": "number"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "array"
+ },
+ {
+ "type": "object"
+ }
+ ]
+ }
+ },
+ "logger_config": {
+ "type": "object",
+ "additionalProperties": {
+ "oneOf": [
+ {
+ "type": "null"
+ },
+ {
+ "type": "boolean"
+ },
+ {
+ "type": "number"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "array"
+ },
+ {
+ "type": "object"
+ }
+ ]
+ }
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "job_uuid",
+ "model",
+ "dataset",
+ "validation_dataset",
+ "algorithm",
+ "algorithm_config",
+ "optimizer_config",
+ "training_config",
+ "hyperparam_search_config",
+ "logger_config"
+ ],
+ "title": "Request to finetune a model."
+ },
+ "LoraFinetuningConfig": {
+ "type": "object",
+ "properties": {
+ "lora_attn_modules": {
+ "type": "array",
+ "items": {
+ "type": "string"
+ }
+ },
+ "apply_lora_to_mlp": {
+ "type": "boolean"
+ },
+ "apply_lora_to_output": {
+ "type": "boolean"
+ },
+ "rank": {
+ "type": "integer"
+ },
+ "alpha": {
+ "type": "integer"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "lora_attn_modules",
+ "apply_lora_to_mlp",
+ "apply_lora_to_output",
+ "rank",
+ "alpha"
+ ]
+ },
+ "OptimizerConfig": {
+ "type": "object",
+ "properties": {
+ "optimizer_type": {
+ "type": "string",
+ "enum": [
+ "adam",
+ "adamw",
+ "sgd"
+ ]
+ },
+ "lr": {
+ "type": "number"
+ },
+ "lr_min": {
+ "type": "number"
+ },
+ "weight_decay": {
+ "type": "number"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "optimizer_type",
+ "lr",
+ "lr_min",
+ "weight_decay"
+ ]
+ },
+ "TrainingConfig": {
+ "type": "object",
+ "properties": {
+ "n_epochs": {
+ "type": "integer"
+ },
+ "batch_size": {
+ "type": "integer"
+ },
+ "shuffle": {
+ "type": "boolean"
+ },
+ "n_iters": {
+ "type": "integer"
+ },
+ "enable_activation_checkpointing": {
+ "type": "boolean"
+ },
+ "memory_efficient_fsdp_wrap": {
+ "type": "boolean"
+ },
+ "fsdp_cpu_offload": {
+ "type": "boolean"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "n_epochs",
+ "batch_size",
+ "shuffle",
+ "n_iters",
+ "enable_activation_checkpointing",
+ "memory_efficient_fsdp_wrap",
+ "fsdp_cpu_offload"
+ ]
}
},
"responses": {}
@@ -1676,17 +2306,23 @@
}
],
"tags": [
+ {
+ "name": "SyntheticDataGeneration"
+ },
{
"name": "RewardScoring"
},
+ {
+ "name": "AgenticSystem"
+ },
+ {
+ "name": "Finetuning"
+ },
{
"name": "Inference"
},
{
- "name": "SyntheticDataGeneration"
- },
- {
- "name": "AgenticSystem"
+ "name": "Datasets"
},
{
"name": "ShieldConfig",
@@ -1728,6 +2364,22 @@
"name": "AgenticSystemExecuteResponseStreamChunk",
"description": "Streamed agent execution response.\n\n"
},
+ {
+ "name": "CreateDatasetRequest",
+ "description": "Request to create a dataset.\n\n"
+ },
+ {
+ "name": "Dataset",
+ "description": "Dataset to be used for training or evaluating language models.\n\n"
+ },
+ {
+ "name": "FinetuningJobStatusResponse",
+ "description": "Status of a finetuning job.\n\n"
+ },
+ {
+ "name": "FinetuningJobLogStream",
+ "description": "Stream of logs from a finetuning job.\n\n"
+ },
{
"name": "ChatCompletionRequest",
"description": ""
@@ -1767,6 +2419,22 @@
{
"name": "RewardScoringResponse",
"description": "Response from the reward scoring. Batch of (prompt, response, score) tuples that pass the threshold.\n\n"
+ },
+ {
+ "name": "FinetuningTrainRequest",
+ "description": "Request to finetune a model.\n\n"
+ },
+ {
+ "name": "LoraFinetuningConfig",
+ "description": ""
+ },
+ {
+ "name": "OptimizerConfig",
+ "description": ""
+ },
+ {
+ "name": "TrainingConfig",
+ "description": ""
}
],
"x-tagGroups": [
@@ -1774,6 +2442,8 @@
"name": "Operations",
"tags": [
"AgenticSystem",
+ "Datasets",
+ "Finetuning",
"Inference",
"RewardScoring",
"SyntheticDataGeneration"
@@ -1795,12 +2465,20 @@
"CompletionRequest",
"CompletionResponse",
"CompletionResponseStreamChunk",
+ "CreateDatasetRequest",
+ "Dataset",
+ "FinetuningJobLogStream",
+ "FinetuningJobStatusResponse",
+ "FinetuningTrainRequest",
+ "LoraFinetuningConfig",
"Message",
+ "OptimizerConfig",
"RewardScoringRequest",
"RewardScoringResponse",
"ShieldConfig",
"SyntheticDataGenerationRequest",
"SyntheticDataGenerationResponse",
+ "TrainingConfig",
"URL"
]
}
diff --git a/source/openapi.yaml b/source/openapi.yaml
index 02e85beed..f8cae85b0 100644
--- a/source/openapi.yaml
+++ b/source/openapi.yaml
@@ -172,13 +172,20 @@ components:
tool_response_delta:
additionalProperties: false
properties:
- response:
- type: string
+ content:
+ oneOf:
+ - type: string
+ - $ref: '#/components/schemas/Attachment'
+ - items:
+ oneOf:
+ - type: string
+ - $ref: '#/components/schemas/Attachment'
+ type: array
tool_name:
type: string
required:
- tool_name
- - response
+ - content
type: object
violation:
additionalProperties: false
@@ -274,13 +281,20 @@ components:
items:
additionalProperties: false
properties:
- response:
- type: string
+ content:
+ oneOf:
+ - type: string
+ - $ref: '#/components/schemas/Attachment'
+ - items:
+ oneOf:
+ - type: string
+ - $ref: '#/components/schemas/Attachment'
+ type: array
tool_name:
type: string
required:
- tool_name
- - response
+ - content
type: object
type: array
uuid:
@@ -690,6 +704,224 @@ components:
- text_delta
title: streamed completion response.
type: object
+ CreateDatasetRequest:
+ additionalProperties: false
+ properties:
+ dataset:
+ $ref: '#/components/schemas/Dataset'
+ uuid:
+ type: string
+ required:
+ - uuid
+ - dataset
+ title: Request to create a dataset.
+ type: object
+ Dataset:
+ additionalProperties: false
+ properties:
+ columns:
+ additionalProperties:
+ enum:
+ - dialog
+ - text
+ - media
+ - number
+ - json
+ type: string
+ type: object
+ content_url:
+ $ref: '#/components/schemas/URL'
+ metadata:
+ additionalProperties:
+ oneOf:
+ - type: 'null'
+ - type: boolean
+ - type: number
+ - type: string
+ - type: array
+ - type: object
+ type: object
+ required:
+ - columns
+ - content_url
+ - metadata
+ title: Dataset to be used for training or evaluating language models.
+ type: object
+ FinetuningJobLogStream:
+ additionalProperties: false
+ properties:
+ job_uuid:
+ type: string
+ log_lines:
+ items:
+ type: string
+ type: array
+ required:
+ - job_uuid
+ - log_lines
+ title: Stream of logs from a finetuning job.
+ type: object
+ FinetuningJobStatusResponse:
+ additionalProperties: false
+ properties:
+ completed_at:
+ format: date-time
+ type: string
+ job_uuid:
+ type: string
+ resources_allocated:
+ additionalProperties:
+ oneOf:
+ - type: 'null'
+ - type: boolean
+ - type: number
+ - type: string
+ - type: array
+ - type: object
+ type: object
+ scheduled_at:
+ format: date-time
+ type: string
+ started_at:
+ format: date-time
+ type: string
+ status:
+ enum:
+ - running
+ - completed
+ - failed
+ - scheduled
+ type: string
+ required:
+ - job_uuid
+ - status
+ title: Status of a finetuning job.
+ type: object
+ FinetuningTrainRequest:
+ additionalProperties: false
+ properties:
+ algorithm:
+ enum:
+ - full
+ - lora
+ - qlora
+ - dora
+ type: string
+ algorithm_config:
+ oneOf:
+ - $ref: '#/components/schemas/LoraFinetuningConfig'
+ - additionalProperties: false
+ properties:
+ alpha:
+ type: integer
+ apply_lora_to_mlp:
+ type: boolean
+ apply_lora_to_output:
+ type: boolean
+ lora_attn_modules:
+ items:
+ type: string
+ type: array
+ rank:
+ type: integer
+ required:
+ - lora_attn_modules
+ - apply_lora_to_mlp
+ - apply_lora_to_output
+ - rank
+ - alpha
+ type: object
+ - additionalProperties: false
+ properties:
+ alpha:
+ type: integer
+ apply_lora_to_mlp:
+ type: boolean
+ apply_lora_to_output:
+ type: boolean
+ lora_attn_modules:
+ items:
+ type: string
+ type: array
+ rank:
+ type: integer
+ required:
+ - lora_attn_modules
+ - apply_lora_to_mlp
+ - apply_lora_to_output
+ - rank
+ - alpha
+ type: object
+ dataset:
+ $ref: '#/components/schemas/Dataset'
+ hyperparam_search_config:
+ additionalProperties:
+ oneOf:
+ - type: 'null'
+ - type: boolean
+ - type: number
+ - type: string
+ - type: array
+ - type: object
+ type: object
+ job_uuid:
+ type: string
+ logger_config:
+ additionalProperties:
+ oneOf:
+ - type: 'null'
+ - type: boolean
+ - type: number
+ - type: string
+ - type: array
+ - type: object
+ type: object
+ model:
+ enum:
+ - llama3_8b
+ - llama3_70b
+ type: string
+ optimizer_config:
+ $ref: '#/components/schemas/OptimizerConfig'
+ training_config:
+ $ref: '#/components/schemas/TrainingConfig'
+ validation_dataset:
+ $ref: '#/components/schemas/Dataset'
+ required:
+ - job_uuid
+ - model
+ - dataset
+ - validation_dataset
+ - algorithm
+ - algorithm_config
+ - optimizer_config
+ - training_config
+ - hyperparam_search_config
+ - logger_config
+ title: Request to finetune a model.
+ type: object
+ LoraFinetuningConfig:
+ additionalProperties: false
+ properties:
+ alpha:
+ type: integer
+ apply_lora_to_mlp:
+ type: boolean
+ apply_lora_to_output:
+ type: boolean
+ lora_attn_modules:
+ items:
+ type: string
+ type: array
+ rank:
+ type: integer
+ required:
+ - lora_attn_modules
+ - apply_lora_to_mlp
+ - apply_lora_to_output
+ - rank
+ - alpha
+ type: object
Message:
additionalProperties: false
properties:
@@ -735,13 +967,20 @@ components:
items:
additionalProperties: false
properties:
- response:
- type: string
+ content:
+ oneOf:
+ - type: string
+ - $ref: '#/components/schemas/Attachment'
+ - items:
+ oneOf:
+ - type: string
+ - $ref: '#/components/schemas/Attachment'
+ type: array
tool_name:
type: string
required:
- tool_name
- - response
+ - content
type: object
type: array
required:
@@ -750,6 +989,27 @@ components:
- tool_calls
- tool_responses
type: object
+ OptimizerConfig:
+ additionalProperties: false
+ properties:
+ lr:
+ type: number
+ lr_min:
+ type: number
+ optimizer_type:
+ enum:
+ - adam
+ - adamw
+ - sgd
+ type: string
+ weight_decay:
+ type: number
+ required:
+ - optimizer_type
+ - lr
+ - lr_min
+ - weight_decay
+ type: object
RewardScoringRequest:
additionalProperties: false
properties:
@@ -885,6 +1145,32 @@ components:
title: Response from the synthetic data generation. Batch of (prompt, response,
score) tuples that pass the threshold.
type: object
+ TrainingConfig:
+ additionalProperties: false
+ properties:
+ batch_size:
+ type: integer
+ enable_activation_checkpointing:
+ type: boolean
+ fsdp_cpu_offload:
+ type: boolean
+ memory_efficient_fsdp_wrap:
+ type: boolean
+ n_epochs:
+ type: integer
+ n_iters:
+ type: integer
+ shuffle:
+ type: boolean
+ required:
+ - n_epochs
+ - batch_size
+ - shuffle
+ - n_iters
+ - enable_activation_checkpointing
+ - memory_efficient_fsdp_wrap
+ - fsdp_cpu_offload
+ type: object
URL:
format: uri
pattern: ^(https?://|file://|data:)
@@ -989,6 +1275,98 @@ paths:
description: Normal completion response. **OR** streamed completion response.
tags:
- Inference
+ /datasets/create:
+ post:
+ parameters: []
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/CreateDatasetRequest'
+ required: true
+ responses:
+ '200':
+ description: OK
+ tags:
+ - Datasets
+ /datasets/delete:
+ delete:
+ parameters:
+ - in: query
+ name: dataset_id
+ required: true
+ schema:
+ type: string
+ responses:
+ '200':
+ description: OK
+ tags:
+ - Datasets
+ /datasets/get:
+ get:
+ parameters:
+ - in: query
+ name: dataset_id
+ required: true
+ schema:
+ type: string
+ responses:
+ '200':
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/Dataset'
+ description: OK
+ tags:
+ - Datasets
+ /finetuning/job/logs:
+ get:
+ parameters:
+ - in: query
+ name: job_uuid
+ required: true
+ schema:
+ type: string
+ responses:
+ '200':
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/FinetuningJobLogStream'
+ description: OK
+ tags:
+ - Finetuning
+ /finetuning/job/status:
+ get:
+ parameters:
+ - in: query
+ name: job_uuid
+ required: true
+ schema:
+ type: string
+ responses:
+ '200':
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/FinetuningJobStatusResponse'
+ description: OK
+ tags:
+ - Finetuning
+ /finetuning/text_generation/train:
+ post:
+ parameters: []
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/FinetuningTrainRequest'
+ required: true
+ responses:
+ '200':
+ description: OK
+ tags:
+ - Finetuning
/reward_scoring/score:
post:
parameters: []
@@ -1030,10 +1408,12 @@ security:
servers:
- url: http://llama.meta.com
tags:
-- name: RewardScoring
-- name: Inference
- name: SyntheticDataGeneration
+- name: RewardScoring
- name: AgenticSystem
+- name: Finetuning
+- name: Inference
+- name: Datasets
- description:
name: ShieldConfig
- description: '
name: AgenticSystemExecuteResponseStreamChunk
+- description: 'Request to create a dataset.
+
+
+ '
+ name: CreateDatasetRequest
+- description: 'Dataset to be used for training or evaluating language models.
+
+
+ '
+ name: Dataset
+- description: 'Status of a finetuning job.
+
+
+ '
+ name: FinetuningJobStatusResponse
+- description: 'Stream of logs from a finetuning job.
+
+
+ '
+ name: FinetuningJobLogStream
- description:
name: ChatCompletionRequest
@@ -1127,10 +1528,25 @@ tags:
'
name: RewardScoringResponse
+- description: 'Request to finetune a model.
+
+
+ '
+ name: FinetuningTrainRequest
+- description:
+ name: LoraFinetuningConfig
+- description:
+ name: OptimizerConfig
+- description:
+ name: TrainingConfig
x-tagGroups:
- name: Operations
tags:
- AgenticSystem
+ - Datasets
+ - Finetuning
- Inference
- RewardScoring
- SyntheticDataGeneration
@@ -1149,10 +1565,18 @@ x-tagGroups:
- CompletionRequest
- CompletionResponse
- CompletionResponseStreamChunk
+ - CreateDatasetRequest
+ - Dataset
+ - FinetuningJobLogStream
+ - FinetuningJobStatusResponse
+ - FinetuningTrainRequest
+ - LoraFinetuningConfig
- Message
+ - OptimizerConfig
- RewardScoringRequest
- RewardScoringResponse
- ShieldConfig
- SyntheticDataGenerationRequest
- SyntheticDataGenerationResponse
+ - TrainingConfig
- URL