kiln-ai 0.8.0__py3-none-any.whl → 0.11.1__py3-none-any.whl

kiln_ai/adapters/fine_tune/dataset_formatter.py

@@ -1,11 +1,13 @@
 import json
 import tempfile
+from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 from typing import Any, Dict, Protocol
 from uuid import uuid4
 
-from kiln_ai.datamodel import DatasetSplit, TaskRun
+from kiln_ai.adapters.model_adapters.base_adapter import COT_FINAL_ANSWER_PROMPT
+from kiln_ai.datamodel import DatasetSplit, FinetuneDataStrategy, TaskRun
 
 
 class DatasetFormat(str, Enum):
@@ -14,6 +16,9 @@ class DatasetFormat(str, Enum):
     """OpenAI chat format with plaintext response"""
     OPENAI_CHAT_JSONL = "openai_chat_jsonl"
 
+    """OpenAI chat format with json response_format"""
+    OPENAI_CHAT_JSON_SCHEMA_JSONL = "openai_chat_json_schema_jsonl"
+
     """OpenAI chat format with tool call response"""
     OPENAI_CHAT_TOOLCALL_JSONL = "openai_chat_toolcall_jsonl"
 
@@ -25,116 +30,338 @@ class DatasetFormat(str, Enum):
         "huggingface_chat_template_toolcall_jsonl"
     )
 
+    """Vertex Gemini 1.5 format (flash and pro)"""
+    VERTEX_GEMINI_1_5 = "vertex_gemini_1_5"
+
+
+@dataclass
+class ModelTrainingData:
+    input: str
+    system_message: str
+    final_output: str
+    # These 3 are optional, and used for COT/Thinking style multi-message responses
+    thinking_instructions: str | None = None
+    thinking: str | None = None
+    thinking_final_answer_prompt: str | None = None
+
+    def supports_cot(self) -> bool:
+        return (
+            self.thinking_instructions is not None
+            and self.thinking is not None
+            and self.thinking_final_answer_prompt is not None
+        )
+
 
 class FormatGenerator(Protocol):
     """Protocol for format generators"""
 
-    def __call__(self, task_run: TaskRun, system_message: str) -> Dict[str, Any]: ...
+    def __call__(
+        self,
+        training_data: ModelTrainingData,
+    ) -> Dict[str, Any]: ...
+
+
+def build_training_data(
+    task_run: TaskRun,
+    system_message: str,
+    include_cot: bool,
+    thinking_instructions: str | None = None,
+) -> ModelTrainingData:
+    """
+    Generate data for training.
+
+    For final output, get the best task output from the task run, preferring repaired output if available.
+
+    For thinking, get the intermediate output if it exists, otherwise return None.
+    """
+    final_output = task_run.output.output
+    if task_run.repaired_output is not None:
+        final_output = task_run.repaired_output.output
+
+    thinking = None
+    thinking_final_answer_prompt = None
+    parent_task = task_run.parent_task()
+
+    if include_cot and task_run.has_thinking_training_data():
+        if not parent_task:
+            raise ValueError(
+                "TaskRuns for training required a parent Task for building a chain of thought prompts. Train without COT, or save this TaskRun to a parent Task."
+            )
+
+        # Prefer reasoning to cot if both are present
+        intermediate_outputs = task_run.intermediate_outputs or {}
+        thinking = intermediate_outputs.get("reasoning") or intermediate_outputs.get(
+            "chain_of_thought"
+        )
+
+        thinking_final_answer_prompt = COT_FINAL_ANSWER_PROMPT
+
+        # Always use the passed thinking instructions, but check they are present for COT
+        if not thinking_instructions:
+            raise ValueError(
+                "Thinking instructions are required when data_strategy is final_and_intermediate"
+            )
+
+    return ModelTrainingData(
+        input=task_run.input,
+        system_message=system_message,
+        final_output=final_output,
+        thinking=thinking,
+        thinking_instructions=thinking_instructions,
+        thinking_final_answer_prompt=thinking_final_answer_prompt,
+    )
 
 
 def generate_chat_message_response(
-    task_run: TaskRun, system_message: str
+    training_data: ModelTrainingData,
 ) -> Dict[str, Any]:
     """Generate OpenAI chat format with plaintext response"""
-    return {
-        "messages": [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": task_run.input},
-            {"role": "assistant", "content": task_run.output.output},
-        ]
-    }
+
+    messages: list[dict[str, str | None]] = [
+        {"role": "system", "content": training_data.system_message},
+        {"role": "user", "content": training_data.input},
+    ]
+
+    if training_data.supports_cot():
+        messages.extend(
+            [
+                {"role": "user", "content": training_data.thinking_instructions},
+                {"role": "assistant", "content": training_data.thinking},
+                {
+                    "role": "user",
+                    "content": training_data.thinking_final_answer_prompt,
+                },
+            ]
+        )
+
+    messages.append({"role": "assistant", "content": training_data.final_output})
+
+    return {"messages": messages}
+
+
+def generate_json_schema_message(
+    training_data: ModelTrainingData,
+) -> Dict[str, Any]:
+    """Generate OpenAI chat format with validated JSON response"""
+    # Load and dump to ensure it's valid JSON and goes to 1 line
+    try:
+        json_data = json.loads(training_data.final_output)
+    except json.JSONDecodeError as e:
+        raise ValueError(
+            f"Invalid JSON in JSON Schema training set: {e}\nOutput Data: {training_data.final_output}"
+        ) from e
+    json_string = json.dumps(json_data, ensure_ascii=False)
+
+    messages: list[dict[str, str | None]] = [
+        {"role": "system", "content": training_data.system_message},
+        {"role": "user", "content": training_data.input},
+    ]
+
+    if training_data.supports_cot():
+        messages.extend(
+            [
+                {"role": "user", "content": training_data.thinking_instructions},
+                {"role": "assistant", "content": training_data.thinking},
+                {
+                    "role": "user",
+                    "content": training_data.thinking_final_answer_prompt,
+                },
+            ]
+        )
+
+    messages.append({"role": "assistant", "content": json_string})
+
+    return {"messages": messages}
 
 
 def generate_chat_message_toolcall(
-    task_run: TaskRun, system_message: str
+    training_data: ModelTrainingData,
 ) -> Dict[str, Any]:
     """Generate OpenAI chat format with tool call response"""
     try:
-        arguments = json.loads(task_run.output.output)
+        arguments = json.loads(training_data.final_output)
     except json.JSONDecodeError as e:
         raise ValueError(f"Invalid JSON in for tool call: {e}") from e
 
-    return {
-        "messages": [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": task_run.input},
-            {
-                "role": "assistant",
-                "content": None,
-                "tool_calls": [
-                    {
-                        "id": "call_1",
-                        "type": "function",
-                        "function": {
-                            "name": "task_response",
-                            # Yes we parse then dump again. This ensures it's valid JSON, and ensures it goes to 1 line
-                            "arguments": json.dumps(arguments),
-                        },
-                    }
-                ],
-            },
-        ]
-    }
+    messages: list[dict[str, Any]] = [
+        {"role": "system", "content": training_data.system_message},
+        {"role": "user", "content": training_data.input},
+    ]
+
+    if training_data.supports_cot():
+        messages.extend(
+            [
+                {"role": "user", "content": training_data.thinking_instructions},
+                {"role": "assistant", "content": training_data.thinking},
+                {
+                    "role": "user",
+                    "content": training_data.thinking_final_answer_prompt,
+                },
+            ]
+        )
+
+    messages.append(
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "call_1",
+                    "type": "function",
+                    "function": {
+                        "name": "task_response",
+                        # Yes we parse then dump again. This ensures it's valid JSON, and ensures it goes to 1 line
+                        "arguments": json.dumps(arguments, ensure_ascii=False),
+                    },
+                }
+            ],
+        },
+    )
+
+    return {"messages": messages}
 
 
 def generate_huggingface_chat_template(
-    task_run: TaskRun, system_message: str
+    training_data: ModelTrainingData,
 ) -> Dict[str, Any]:
     """Generate HuggingFace chat template"""
-    return {
-        "conversations": [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": task_run.input},
-            {"role": "assistant", "content": task_run.output.output},
-        ]
-    }
+
+    conversations: list[dict[str, Any]] = [
+        {"role": "system", "content": training_data.system_message},
+        {"role": "user", "content": training_data.input},
+    ]
+
+    if training_data.supports_cot():
+        conversations.extend(
+            [
+                {"role": "user", "content": training_data.thinking_instructions},
+                {"role": "assistant", "content": training_data.thinking},
+                {"role": "user", "content": training_data.thinking_final_answer_prompt},
+            ]
+        )
+
+    conversations.append({"role": "assistant", "content": training_data.final_output})
+
+    return {"conversations": conversations}
 
 
 def generate_huggingface_chat_template_toolcall(
-    task_run: TaskRun, system_message: str
+    training_data: ModelTrainingData,
 ) -> Dict[str, Any]:
     """Generate HuggingFace chat template with tool calls"""
     try:
-        arguments = json.loads(task_run.output.output)
+        arguments = json.loads(training_data.final_output)
     except json.JSONDecodeError as e:
         raise ValueError(f"Invalid JSON in for tool call: {e}") from e
 
     # See https://huggingface.co/docs/transformers/en/chat_templating
+    conversations: list[dict[str, Any]] = [
+        {"role": "system", "content": training_data.system_message},
+        {"role": "user", "content": training_data.input},
+    ]
+
+    if training_data.supports_cot():
+        conversations.extend(
+            [
+                {"role": "user", "content": training_data.thinking_instructions},
+                {"role": "assistant", "content": training_data.thinking},
+                {"role": "user", "content": training_data.thinking_final_answer_prompt},
+            ]
+        )
+
+    conversations.append(
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "task_response",
+                        "id": str(uuid4()).replace("-", "")[:9],
+                        "arguments": arguments,
+                    },
+                }
+            ],
+        },
+    )
+
+    return {"conversations": conversations}
+
+
+def generate_vertex_gemini_1_5(
+    training_data: ModelTrainingData,
+) -> Dict[str, Any]:
+    """Generate Vertex Gemini 1.5 format (flash and pro)"""
+    # See https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini-supervised-tuning-prepare
+
+    contents = [
+        {
+            "role": "user",
+            "parts": [
+                {
+                    "text": training_data.input,
+                }
+            ],
+        }
+    ]
+
+    if training_data.supports_cot():
+        contents.extend(
+            [
+                {
+                    "role": "user",
+                    "parts": [{"text": training_data.thinking_instructions}],
+                },
+                {"role": "model", "parts": [{"text": training_data.thinking}]},
+                {
+                    "role": "user",
+                    "parts": [{"text": training_data.thinking_final_answer_prompt}],
+                },
+            ]
+        )
+
+    contents.append(
+        {
+            "role": "model",
+            "parts": [{"text": training_data.final_output}],
+        }
+    )
+
     return {
-        "conversations": [
-            {"role": "system", "content": system_message},
-            {"role": "user", "content": task_run.input},
-            {
-                "role": "assistant",
-                "tool_calls": [
-                    {
-                        "type": "function",
-                        "function": {
-                            "name": "task_response",
-                            "id": str(uuid4()).replace("-", "")[:9],
-                            "arguments": arguments,
-                        },
-                    }
-                ],
-            },
-        ]
+        "systemInstruction": {
+            "role": "system",
+            "parts": [
+                {
+                    "text": training_data.system_message,
+                }
+            ],
+        },
+        "contents": contents,
     }
 
 
 FORMAT_GENERATORS: Dict[DatasetFormat, FormatGenerator] = {
     DatasetFormat.OPENAI_CHAT_JSONL: generate_chat_message_response,
+    DatasetFormat.OPENAI_CHAT_JSON_SCHEMA_JSONL: generate_json_schema_message,
     DatasetFormat.OPENAI_CHAT_TOOLCALL_JSONL: generate_chat_message_toolcall,
     DatasetFormat.HUGGINGFACE_CHAT_TEMPLATE_JSONL: generate_huggingface_chat_template,
     DatasetFormat.HUGGINGFACE_CHAT_TEMPLATE_TOOLCALL_JSONL: generate_huggingface_chat_template_toolcall,
+    DatasetFormat.VERTEX_GEMINI_1_5: generate_vertex_gemini_1_5,
 }
 
 
 class DatasetFormatter:
     """Handles formatting of datasets into various output formats"""
 
-    def __init__(self, dataset: DatasetSplit, system_message: str):
+    def __init__(
+        self,
+        dataset: DatasetSplit,
+        system_message: str,
+        thinking_instructions: str | None = None,
+    ):
         self.dataset = dataset
         self.system_message = system_message
+        self.thinking_instructions = thinking_instructions
 
         task = dataset.parent_task()
         if task is None:
@@ -142,7 +369,11 @@ class DatasetFormatter:
         self.task = task
 
     def dump_to_file(
-        self, split_name: str, format_type: DatasetFormat, path: Path | None = None
+        self,
+        split_name: str,
+        format_type: DatasetFormat,
+        data_strategy: FinetuneDataStrategy,
+        path: Path | None = None,
     ) -> Path:
         """
         Format the dataset into the specified format.
@@ -154,6 +385,10 @@ class DatasetFormatter:
 
         Returns:
             Path to the generated file
+
+        Note:
+            The output is written in UTF-8 encoding with ensure_ascii=False to properly
+            support international text content while maintaining readability.
         """
         if format_type not in FORMAT_GENERATORS:
             raise ValueError(f"Unsupported format: {format_type}")
@@ -162,11 +397,13 @@ class DatasetFormatter:
 
         generator = FORMAT_GENERATORS[format_type]
 
+        include_cot = data_strategy == FinetuneDataStrategy.final_and_intermediate
+
         # Write to a temp file if no path is provided
         output_path = (
             path
             or Path(tempfile.gettempdir())
-            / f"{self.dataset.name}_{split_name}_{format_type}.jsonl"
+            / f"{self.dataset.name} -- split-{split_name} -- format-{format_type.value} -- {'cot' if include_cot else 'no-cot'}.jsonl"
         )
 
         runs = self.task.runs()
@@ -181,7 +418,15 @@ class DatasetFormatter:
                         f"Task run {run_id} not found. This is required by this dataset."
                     )
 
-                example = generator(task_run, self.system_message)
-                f.write(json.dumps(example) + "\n")
+                training_data = build_training_data(
+                    task_run=task_run,
+                    system_message=self.system_message,
+                    include_cot=include_cot,
+                    thinking_instructions=self.thinking_instructions,
+                )
+                example = generator(training_data)
+                # Allow non-ascii characters in the dataset.
+                # Better readability for non-English users. If you don't support UTF-8... you should.
+                f.write(json.dumps(example, ensure_ascii=False) + "\n")
 
         return output_path

kiln_ai/adapters/fine_tune/fireworks_finetune.py

@@ -1,3 +1,4 @@
+from typing import Tuple
 from uuid import uuid4
 
 import httpx
@@ -9,7 +10,7 @@ from kiln_ai.adapters.fine_tune.base_finetune import (
     FineTuneStatusType,
 )
 from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetFormatter
-from kiln_ai.datamodel import DatasetSplit, Task
+from kiln_ai.datamodel import DatasetSplit, StructuredOutputMode, Task
 from kiln_ai.utils.config import Config
 
 
@@ -19,7 +20,7 @@ class FireworksFinetune(BaseFinetuneAdapter):
     """
 
     async def status(self) -> FineTuneStatus:
-        status = await self._status()
+        status, _ = await self._status()
         # update the datamodel if the status has changed
         if self.datamodel.latest_status != status.status:
             self.datamodel.latest_status = status.status
@@ -34,7 +35,7 @@ class FireworksFinetune(BaseFinetuneAdapter):
 
         return status
 
-    async def _status(self) -> FineTuneStatus:
+    async def _status(self) -> Tuple[FineTuneStatus, str | None]:
         try:
             api_key = Config.shared().fireworks_api_key
             account_id = Config.shared().fireworks_account_id
@@ -42,13 +43,13 @@ class FireworksFinetune(BaseFinetuneAdapter):
                 return FineTuneStatus(
                     status=FineTuneStatusType.unknown,
                     message="Fireworks API key or account ID not set",
-                )
+                ), None
             fine_tuning_job_id = self.datamodel.provider_id
             if not fine_tuning_job_id:
                 return FineTuneStatus(
                     status=FineTuneStatusType.unknown,
                     message="Fine-tuning job ID not set. Can not retrieve status.",
-                )
+                ), None
             # Fireworks uses path style IDs
             url = f"https://api.fireworks.ai/v1/{fine_tuning_job_id}"
             headers = {"Authorization": f"Bearer {api_key}"}
@@ -60,49 +61,63 @@ class FireworksFinetune(BaseFinetuneAdapter):
                 return FineTuneStatus(
                     status=FineTuneStatusType.unknown,
                     message=f"Error retrieving fine-tuning job status: [{response.status_code}] {response.text}",
-                )
+                ), None
             data = response.json()
+            model_id = data.get("outputModel")
 
             if "state" not in data:
                 return FineTuneStatus(
                     status=FineTuneStatusType.unknown,
                     message="Invalid response from Fireworks (no state).",
-                )
+                ), model_id
 
             state = data["state"]
-            if state in ["FAILED", "DELETING"]:
+            if state in ["FAILED", "DELETING", "JOB_STATE_FAILED"]:
                 return FineTuneStatus(
                     status=FineTuneStatusType.failed,
                     message="Fine-tuning job failed",
-                )
-            elif state in ["CREATING", "PENDING", "RUNNING"]:
+                ), model_id
+            elif state in [
+                "CREATING",
+                "PENDING",
+                "RUNNING",
+                "JOB_STATE_VALIDATING",
+                "JOB_STATE_RUNNING",
+            ]:
                 return FineTuneStatus(
                     status=FineTuneStatusType.running,
                     message=f"Fine-tuning job is running [{state}]",
-                )
-            elif state == "COMPLETED":
+                ), model_id
+            elif state in ["COMPLETED", "JOB_STATE_COMPLETED"]:
                 return FineTuneStatus(
                     status=FineTuneStatusType.completed,
                     message="Fine-tuning job completed",
-                )
+                ), model_id
             else:
                 return FineTuneStatus(
                     status=FineTuneStatusType.unknown,
                     message=f"Unknown fine-tuning job status [{state}]",
-                )
+                ), model_id
         except Exception as e:
             return FineTuneStatus(
                 status=FineTuneStatusType.unknown,
                 message=f"Error retrieving fine-tuning job status: {e}",
-            )
+            ), None
 
     async def _start(self, dataset: DatasetSplit) -> None:
         task = self.datamodel.parent_task()
         if not task:
             raise ValueError("Task is required to start a fine-tune")
 
+        format = DatasetFormat.OPENAI_CHAT_JSONL
+        if task.output_json_schema:
+            # This formatter will check it's valid JSON, and normalize the output (chat format just uses exact string).
+            format = DatasetFormat.OPENAI_CHAT_JSON_SCHEMA_JSONL
+            # Fireworks doesn't support function calls or json schema, so we'll use json mode at call time
+            self.datamodel.structured_output_mode = StructuredOutputMode.json_mode
+
         train_file_id = await self.generate_and_upload_jsonl(
-            dataset, self.datamodel.train_split_name, task
+            dataset, self.datamodel.train_split_name, task, format
         )
 
         api_key = Config.shared().fireworks_api_key
@@ -110,9 +125,7 @@ class FireworksFinetune(BaseFinetuneAdapter):
         if not api_key or not account_id:
             raise ValueError("Fireworks API key or account ID not set")
 
-        url = f"https://api.fireworks.ai/v1/accounts/{account_id}/fineTuningJobs"
-        # Model ID != fine tune ID on Fireworks. Model is the result of the tune job.
-        model_id = str(uuid4())
+        url = f"https://api.fireworks.ai/v1/accounts/{account_id}/supervisedFineTuningJobs"
         # Limit the display name to 60 characters
         display_name = (
             f"Kiln AI fine-tuning [ID:{self.datamodel.id}][name:{self.datamodel.name}]"[
@@ -120,11 +133,9 @@ class FireworksFinetune(BaseFinetuneAdapter):
             ]
         )
         payload = {
-            "modelId": model_id,
             "dataset": f"accounts/{account_id}/datasets/{train_file_id}",
             "displayName": display_name,
             "baseModel": self.datamodel.base_model_id,
-            "conversation": {},
         }
         hyperparameters = self.create_payload_parameters(self.datamodel.parameters)
         payload.update(hyperparameters)
@@ -148,21 +159,22 @@ class FireworksFinetune(BaseFinetuneAdapter):
         # model ID is the model that results from the fine-tune job
         job_id = data["name"]
         self.datamodel.provider_id = job_id
-        # Keep track of the expected model ID before it's deployed as a property. We move it to fine_tune_model_id after deployment.
-        self.datamodel.properties["undeployed_model_id"] = (
-            f"accounts/{account_id}/models/{model_id}"
-        )
+
+        # Fireworks has 2 different fine tuning endpoints, and depending which you use, the URLs change
+        self.datamodel.properties["endpoint_version"] = "v2"
+
         if self.datamodel.path:
             self.datamodel.save_to_file()
 
     async def generate_and_upload_jsonl(
-        self, dataset: DatasetSplit, split_name: str, task: Task
+        self, dataset: DatasetSplit, split_name: str, task: Task, format: DatasetFormat
     ) -> str:
-        formatter = DatasetFormatter(dataset, self.datamodel.system_message)
-        # OpenAI compatible: https://docs.fireworks.ai/fine-tuning/fine-tuning-models#conversation
-        # Note: Fireworks does not support tool calls (confirmed by Fireworks team) so we'll use json mode
-        format = DatasetFormat.OPENAI_CHAT_JSONL
-        path = formatter.dump_to_file(split_name, format)
+        formatter = DatasetFormatter(
+            dataset=dataset,
+            system_message=self.datamodel.system_message,
+            thinking_instructions=self.datamodel.thinking_instructions,
+        )
+        path = formatter.dump_to_file(split_name, format, self.datamodel.data_strategy)
 
         # First call creates the dataset
         api_key = Config.shared().fireworks_api_key
@@ -276,7 +288,10 @@ class FireworksFinetune(BaseFinetuneAdapter):
         if not api_key or not account_id:
             raise ValueError("Fireworks API key or account ID not set")
 
-        model_id = self.datamodel.properties.get("undeployed_model_id")
+        # Model ID != fine tune ID on Fireworks. Model is the result of the tune job. Call status to get it.
+        status, model_id = await self._status()
+        if status.status != FineTuneStatusType.completed:
+            return False
         if not model_id or not isinstance(model_id, str):
             return False