kiln-ai 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

kiln_ai/adapters/__init__.py

@@ -12,7 +12,15 @@ The prompt_builders submodule contains classes that build prompts for use with t
 The repair submodule contains an adapter for the repair task.
 """
 
-from . import base_adapter, langchain_adapters, ml_model_list, prompt_builders, repair
+from . import (
+    base_adapter,
+    data_gen,
+    fine_tune,
+    langchain_adapters,
+    ml_model_list,
+    prompt_builders,
+    repair,
+)
 
 __all__ = [
     "base_adapter",
@@ -20,4 +28,6 @@ __all__ = [
     "ml_model_list",
     "prompt_builders",
     "repair",
+    "data_gen",
+    "fine_tune",
 ]

kiln_ai/adapters/adapter_registry.py

@@ -0,0 +1,19 @@
+from kiln_ai import datamodel
+from kiln_ai.adapters.base_adapter import BaseAdapter
+from kiln_ai.adapters.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.prompt_builders import BasePromptBuilder
+
+
+def adapter_for_task(
+    kiln_task: datamodel.Task,
+    model_name: str | None = None,
+    provider: str | None = None,
+    prompt_builder: BasePromptBuilder | None = None,
+) -> BaseAdapter:
+    # We use langchain for everything right now, but can add any others here
+    return LangchainAdapter(
+        kiln_task,
+        model_name=model_name,
+        provider=provider,
+        prompt_builder=prompt_builder,
+    )

kiln_ai/adapters/data_gen/__init__.py

@@ -0,0 +1,11 @@
+"""
+# Data Generation
+
+A task to generate synthetic data for Kiln Tasks. This generates the inputs, which then can be run through the task.
+
+Optional human guidance can be provided to guide the generation process.
+"""
+
+from . import data_gen_task
+
+__all__ = ["data_gen_task"]

kiln_ai/adapters/data_gen/data_gen_task.py

@@ -1,8 +1,9 @@
 import json
 
+from pydantic import BaseModel
+
 from kiln_ai.adapters.prompt_builders import SimplePromptBuilder
 from kiln_ai.datamodel import Project, Task
-from pydantic import BaseModel
 
 from .data_gen_prompts import (
     SAMPLE_GENERATION_PROMPT,
@@ -11,6 +12,16 @@ from .data_gen_prompts import (
 
 
 class DataGenCategoriesTaskInput(BaseModel):
+    """Input model for generating categories/subtopics.
+
+    Attributes:
+        node_path: List of strings representing the hierarchical path to current node
+        system_prompt: System prompt to guide the AI generation
+        num_subtopics: Number of subtopics to generate
+        human_guidance: Optional human guidance to influence generation
+        existing_topics: Optional list of existing topics to avoid duplication
+    """
+
     node_path: list[str]
     system_prompt: str
     num_subtopics: int
@@ -26,6 +37,18 @@ class DataGenCategoriesTaskInput(BaseModel):
         human_guidance: str | None = None,
         existing_topics: list[str] | None = None,
     ) -> "DataGenCategoriesTaskInput":
+        """Create a DataGenCategoriesTaskInput instance from a Task.
+
+        Args:
+            task: The source Task object
+            node_path: Path to current node in topic hierarchy
+            num_subtopics: Number of subtopics to generate
+            human_guidance: Optional guidance for generation
+            existing_topics: Optional list of existing topics
+
+        Returns:
+            A new DataGenCategoriesTaskInput instance
+        """
         prompt_builder = SimplePromptBuilder(task=task)
         return cls(
             node_path=node_path,
@@ -37,10 +60,22 @@ class DataGenCategoriesTaskInput(BaseModel):
 
 
 class DataGenCategoriesTaskOutput(BaseModel):
+    """Output model for generated categories/subtopics.
+
+    Attributes:
+        subtopics: List of generated subtopic strings
+    """
+
     subtopics: list[str]
 
 
 class DataGenCategoriesTask(Task, parent_of={}):
+    """Task for generating hierarchical categories/subtopics.
+
+    Generates synthetic data categories which can be used to generate
+    training data for model learning.
+    """
+
     def __init__(self):
         # Keep the typechecker happy. TODO: shouldn't need this or parent_of above.
         tmp_project = Project(name="DataGen")
@@ -59,6 +94,15 @@ class DataGenCategoriesTask(Task, parent_of={}):
 
 
 class DataGenSampleTaskInput(BaseModel):
+    """Input model for generating data samples for a kiln task.
+
+    Attributes:
+        topic: List of strings representing the topic path
+        system_prompt: System prompt to guide the AI generation
+        num_samples: Number of samples to generate
+        human_guidance: Optional human guidance to influence generation
+    """
+
     topic: list[str]
     system_prompt: str
     num_samples: int
@@ -72,6 +116,17 @@ class DataGenSampleTaskInput(BaseModel):
         num_samples: int = 8,
         human_guidance: str | None = None,
     ) -> "DataGenSampleTaskInput":
+        """Create a DataGenSampleTaskInput instance from a Task.
+
+        Args:
+            task: The source Task object
+            topic: Topic path for sample generation
+            num_samples: Number of samples to generate
+            human_guidance: Optional guidance for generation
+
+        Returns:
+            A new DataGenSampleTaskInput instance
+        """
         prompt_builder = SimplePromptBuilder(task=task)
         return cls(
             topic=topic,
@@ -82,6 +137,14 @@ class DataGenSampleTaskInput(BaseModel):
 
 
 def list_json_schema_for_task(task: Task) -> str:
+    """Generate a JSON schema for a list of task inputs (json schema)
+
+    Args:
+        task: Task object whose input schema will be used
+
+    Returns:
+        JSON string representing the schema for a list of task inputs
+    """
     if task.input_json_schema:
         items_schema = json.loads(task.input_json_schema)
     else:
@@ -104,6 +167,11 @@ def list_json_schema_for_task(task: Task) -> str:
 
 
 class DataGenSampleTask(Task, parent_of={}):
+    """Task for generating data samples for a given topic.
+
+    Generates synthetic data samples based on provided topics and subtopics.
+    """
+
     def __init__(self, target_task: Task, num_samples: int = 8):
         # Keep the typechecker happy. TODO: shouldn't need this or parent_of above.
         tmp_project = Project(name="DataGenSample")

kiln_ai/adapters/data_gen/test_data_gen_task.py

@@ -1,6 +1,8 @@
 import json
 
 import pytest
+
+from kiln_ai.adapters.adapter_registry import adapter_for_task
 from kiln_ai.adapters.data_gen.data_gen_task import (
     DataGenCategoriesTask,
     DataGenCategoriesTaskInput,
@@ -9,8 +11,7 @@ from kiln_ai.adapters.data_gen.data_gen_task import (
     DataGenSampleTaskInput,
     list_json_schema_for_task,
 )
-from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
-from kiln_ai.adapters.ml_model_list import get_model_and_provider
+from kiln_ai.adapters.provider_tools import get_model_and_provider
 from kiln_ai.adapters.test_prompt_adaptors import get_all_models_and_providers
 from kiln_ai.datamodel import Project, Task
 
@@ -107,7 +108,7 @@ async def test_data_gen_all_models_providers(
     data_gen_task = DataGenCategoriesTask()
     data_gen_input = DataGenCategoriesTaskInput.from_task(base_task, num_subtopics=6)
 
-    adapter = LangChainPromptAdapter(
+    adapter = adapter_for_task(
         data_gen_task,
         model_name=model_name,
         provider=provider_name,
@@ -231,7 +232,7 @@ async def test_data_gen_sample_all_models_providers(
         base_task, topic=["riding horses"], num_samples=4
     )
 
-    adapter = LangChainPromptAdapter(
+    adapter = adapter_for_task(
         data_gen_task,
         model_name=model_name,
         provider=provider_name,
@@ -250,17 +251,25 @@ async def test_data_gen_sample_all_models_providers(
 @pytest.mark.ollama
 @pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
 async def test_data_gen_sample_all_models_providers_with_structured_output(
-    tmp_path, model_name, provider_name, base_task
+    tmp_path, model_name, provider_name
 ):
-    base_task.output_json_schema = json.dumps(
-        {
-            "type": "object",
-            "properties": {
-                "opening": {"type": "string"},
-                "closing": {"type": "string"},
-            },
-            "required": ["opening", "closing"],
-        }
+    project = Project(name="TestProject")
+    task = Task(
+        name="Summarize",
+        parent=project,
+        description="Explain if the username matches the tweet",
+        instruction="Explain if the username matches the tweet",
+        requirements=[],
+        input_json_schema=json.dumps(
+            {
+                "type": "object",
+                "properties": {
+                    "username": {"type": "string"},
+                    "tweet": {"type": "string"},
+                },
+                "required": ["username", "tweet"],
+            }
+        ),
     )
 
     _, provider = get_model_and_provider(model_name, provider_name)
@@ -268,12 +277,12 @@ async def test_data_gen_sample_all_models_providers_with_structured_output(
         # pass if the model doesn't support data gen (testing the support flag is part of this)
         return
 
-    data_gen_task = DataGenSampleTask(target_task=base_task)
+    data_gen_task = DataGenSampleTask(target_task=task)
     data_gen_input = DataGenSampleTaskInput.from_task(
-        base_task, topic=["riding horses"], num_samples=4
+        task, topic=["Food"], num_samples=4
     )
 
-    adapter = LangChainPromptAdapter(
+    adapter = adapter_for_task(
         data_gen_task,
         model_name=model_name,
         provider=provider_name,
@@ -286,7 +295,7 @@ async def test_data_gen_sample_all_models_providers_with_structured_output(
     assert len(samples) == 4
     for sample in samples:
         assert isinstance(sample, dict)
-        assert "opening" in sample
-        assert "closing" in sample
-        assert isinstance(sample["opening"], str)
-        assert isinstance(sample["closing"], str)
+        assert "username" in sample
+        assert "tweet" in sample
+        assert isinstance(sample["username"], str)
+        assert isinstance(sample["tweet"], str)

kiln_ai/adapters/fine_tune/__init__.py

@@ -0,0 +1,14 @@
+"""
+# Fine-Tuning
+
+A set of classes for fine-tuning models.
+"""
+
+from . import base_finetune, dataset_formatter, finetune_registry, openai_finetune
+
+__all__ = [
+    "base_finetune",
+    "openai_finetune",
+    "dataset_formatter",
+    "finetune_registry",
+]

kiln_ai/adapters/fine_tune/base_finetune.py

@@ -0,0 +1,186 @@
+from abc import ABC, abstractmethod
+from typing import Literal
+
+from pydantic import BaseModel
+
+from kiln_ai.adapters.ml_model_list import built_in_models
+from kiln_ai.datamodel import DatasetSplit, FineTuneStatusType
+from kiln_ai.datamodel import Finetune as FinetuneModel
+from kiln_ai.utils.name_generator import generate_memorable_name
+
+
+class FineTuneStatus(BaseModel):
+    """
+    The status of a fine-tune, including a user friendly message.
+    """
+
+    status: FineTuneStatusType
+    message: str | None = None
+
+
+class FineTuneParameter(BaseModel):
+    """
+    A parameter for a fine-tune. Hyperparameters, etc.
+    """
+
+    name: str
+    type: Literal["string", "int", "float", "bool"]
+    description: str
+    optional: bool = True
+
+
+TYPE_MAP = {
+    "string": str,
+    "int": int,
+    "float": float,
+    "bool": bool,
+}
+
+
+class BaseFinetuneAdapter(ABC):
+    """
+    A base class for fine-tuning adapters.
+    """
+
+    def __init__(
+        self,
+        datamodel: FinetuneModel,
+    ):
+        self.datamodel = datamodel
+
+    @classmethod
+    async def create_and_start(
+        cls,
+        dataset: DatasetSplit,
+        provider_id: str,
+        provider_base_model_id: str,
+        train_split_name: str,
+        system_message: str,
+        parameters: dict[str, str | int | float | bool] = {},
+        name: str | None = None,
+        description: str | None = None,
+        validation_split_name: str | None = None,
+    ) -> tuple["BaseFinetuneAdapter", FinetuneModel]:
+        """
+        Create and start a fine-tune.
+        """
+
+        cls.check_valid_provider_model(provider_id, provider_base_model_id)
+
+        if not dataset.id:
+            raise ValueError("Dataset must have an id")
+
+        if train_split_name not in dataset.split_contents:
+            raise ValueError(f"Train split {train_split_name} not found in dataset")
+
+        if (
+            validation_split_name
+            and validation_split_name not in dataset.split_contents
+        ):
+            raise ValueError(
+                f"Validation split {validation_split_name} not found in dataset"
+            )
+
+        # Default name if not provided
+        if name is None:
+            name = generate_memorable_name()
+
+        cls.validate_parameters(parameters)
+        parent_task = dataset.parent_task()
+        if parent_task is None or not parent_task.path:
+            raise ValueError("Dataset must have a parent task with a path")
+
+        datamodel = FinetuneModel(
+            name=name,
+            description=description,
+            provider=provider_id,
+            base_model_id=provider_base_model_id,
+            dataset_split_id=dataset.id,
+            train_split_name=train_split_name,
+            validation_split_name=validation_split_name,
+            parameters=parameters,
+            system_message=system_message,
+            parent=parent_task,
+        )
+
+        adapter = cls(datamodel)
+        await adapter._start(dataset)
+
+        datamodel.save_to_file()
+
+        return adapter, datamodel
+
+    @abstractmethod
+    async def _start(self, dataset: DatasetSplit) -> None:
+        """
+        Start the fine-tune.
+        """
+        pass
+
+    @abstractmethod
+    async def status(self) -> FineTuneStatus:
+        """
+        Get the status of the fine-tune.
+        """
+        pass
+
+    @classmethod
+    def available_parameters(cls) -> list[FineTuneParameter]:
+        """
+        Returns a list of parameters that can be provided for this fine-tune. Includes hyperparameters, etc.
+        """
+        return []
+
+    @classmethod
+    def validate_parameters(
+        cls, parameters: dict[str, str | int | float | bool]
+    ) -> None:
+        """
+        Validate the parameters for this fine-tune.
+        """
+        # Check required parameters and parameter types
+        available_parameters = cls.available_parameters()
+        for parameter in available_parameters:
+            if not parameter.optional and parameter.name not in parameters:
+                raise ValueError(f"Parameter {parameter.name} is required")
+            elif parameter.name in parameters:
+                # check parameter is correct type
+                expected_type = TYPE_MAP[parameter.type]
+                value = parameters[parameter.name]
+
+                # Strict type checking for numeric types
+                if expected_type is float and not isinstance(value, float):
+                    raise ValueError(
+                        f"Parameter {parameter.name} must be a float, got {type(value)}"
+                    )
+                elif expected_type is int and not isinstance(value, int):
+                    raise ValueError(
+                        f"Parameter {parameter.name} must be an integer, got {type(value)}"
+                    )
+                elif not isinstance(value, expected_type):
+                    raise ValueError(
+                        f"Parameter {parameter.name} must be type {expected_type}, got {type(value)}"
+                    )
+
+        allowed_parameters = [p.name for p in available_parameters]
+        for parameter_key in parameters:
+            if parameter_key not in allowed_parameters:
+                raise ValueError(f"Parameter {parameter_key} is not available")
+
+    @classmethod
+    def check_valid_provider_model(
+        cls, provider_id: str, provider_base_model_id: str
+    ) -> None:
+        """
+        Check if the provider and base model are valid.
+        """
+        for model in built_in_models:
+            for provider in model.providers:
+                if (
+                    provider.name == provider_id
+                    and provider.provider_finetune_id == provider_base_model_id
+                ):
+                    return
+        raise ValueError(
+            f"Provider {provider_id} with base model {provider_base_model_id} is not available"
+        )

kiln_ai/adapters/fine_tune/dataset_formatter.py

@@ -0,0 +1,187 @@
+import json
+import tempfile
+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
+
+
+class DatasetFormat(str, Enum):
+    """Formats for dataset generation. Both for file format (like JSONL), and internal structure (like chat/toolcall)"""
+
+    """OpenAI chat format with plaintext response"""
+    OPENAI_CHAT_JSONL = "openai_chat_jsonl"
+
+    """OpenAI chat format with tool call response"""
+    OPENAI_CHAT_TOOLCALL_JSONL = "openai_chat_toolcall_jsonl"
+
+    """HuggingFace chat template in JSONL"""
+    HUGGINGFACE_CHAT_TEMPLATE_JSONL = "huggingface_chat_template_jsonl"
+
+    """HuggingFace chat template with tool calls in JSONL"""
+    HUGGINGFACE_CHAT_TEMPLATE_TOOLCALL_JSONL = (
+        "huggingface_chat_template_toolcall_jsonl"
+    )
+
+
+class FormatGenerator(Protocol):
+    """Protocol for format generators"""
+
+    def __call__(self, task_run: TaskRun, system_message: str) -> Dict[str, Any]: ...
+
+
+def generate_chat_message_response(
+    task_run: TaskRun, system_message: str
+) -> 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},
+        ]
+    }
+
+
+def generate_chat_message_toolcall(
+    task_run: TaskRun, system_message: str
+) -> Dict[str, Any]:
+    """Generate OpenAI chat format with tool call response"""
+    try:
+        arguments = json.loads(task_run.output.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),
+                        },
+                    }
+                ],
+            },
+        ]
+    }
+
+
+def generate_huggingface_chat_template(
+    task_run: TaskRun, system_message: str
+) -> 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},
+        ]
+    }
+
+
+def generate_huggingface_chat_template_toolcall(
+    task_run: TaskRun, system_message: str
+) -> Dict[str, Any]:
+    """Generate HuggingFace chat template with tool calls"""
+    try:
+        arguments = json.loads(task_run.output.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
+    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,
+                        },
+                    }
+                ],
+            },
+        ]
+    }
+
+
+FORMAT_GENERATORS: Dict[DatasetFormat, FormatGenerator] = {
+    DatasetFormat.OPENAI_CHAT_JSONL: generate_chat_message_response,
+    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,
+}
+
+
+class DatasetFormatter:
+    """Handles formatting of datasets into various output formats"""
+
+    def __init__(self, dataset: DatasetSplit, system_message: str):
+        self.dataset = dataset
+        self.system_message = system_message
+
+        task = dataset.parent_task()
+        if task is None:
+            raise ValueError("Dataset has no parent task")
+        self.task = task
+
+    def dump_to_file(
+        self, split_name: str, format_type: DatasetFormat, path: Path | None = None
+    ) -> Path:
+        """
+        Format the dataset into the specified format.
+
+        Args:
+            split_name: Name of the split to dump
+            format_type: Format to generate the dataset in
+            path: Optional path to write to. If None, writes to temp directory
+
+        Returns:
+            Path to the generated file
+        """
+        if format_type not in FORMAT_GENERATORS:
+            raise ValueError(f"Unsupported format: {format_type}")
+        if split_name not in self.dataset.split_contents:
+            raise ValueError(f"Split {split_name} not found in dataset")
+
+        generator = FORMAT_GENERATORS[format_type]
+
+        # 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"
+        )
+
+        runs = self.task.runs()
+        runs_by_id = {run.id: run for run in runs}
+
+        # Generate formatted output with UTF-8 encoding
+        with open(output_path, "w", encoding="utf-8") as f:
+            for run_id in self.dataset.split_contents[split_name]:
+                task_run = runs_by_id[run_id]
+                if task_run is None:
+                    raise ValueError(
+                        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")
+
+        return output_path

kiln_ai/adapters/fine_tune/finetune_registry.py

@@ -0,0 +1,11 @@
+from typing import Type
+
+from kiln_ai.adapters.fine_tune.base_finetune import BaseFinetuneAdapter
+from kiln_ai.adapters.fine_tune.fireworks_finetune import FireworksFinetune
+from kiln_ai.adapters.fine_tune.openai_finetune import OpenAIFinetune
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+
+finetune_registry: dict[ModelProviderName, Type[BaseFinetuneAdapter]] = {
+    ModelProviderName.openai: OpenAIFinetune,
+    ModelProviderName.fireworks_ai: FireworksFinetune,
+}