From 69ecf55de27b2f10777ec2d78177158f13d2e045 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Wed, 10 Jul 2024 20:47:05 -0700 Subject: [PATCH] finetuning --- source/api_definitions.py | 110 +++++- source/finetuning_types.py | 98 ++++++ source/model_types.py | 4 +- source/openapi.html | 704 ++++++++++++++++++++++++++++++++++++- source/openapi.yaml | 446 ++++++++++++++++++++++- 5 files changed, 1334 insertions(+), 28 deletions(-) create mode 100644 source/finetuning_types.py 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