kiln-ai 0.8.1__py3-none-any.whl → 0.12.0__py3-none-any.whl

kiln_ai/datamodel/basemodel.py

@@ -120,11 +120,12 @@ class KilnBaseModel(BaseModel):
         return cls.load_from_file(path)
 
     @classmethod
-    def load_from_file(cls: Type[T], path: Path | str) -> T:
+    def load_from_file(cls: Type[T], path: Path | str, readonly: bool = False) -> T:
         """Load a model instance from a specific file path.
 
         Args:
             path (Path): Path to the model file
+            readonly (bool): If True, the model will be returned in readonly mode (cached instance, not a copy, not safe to mutate)
 
         Returns:
             T: Instance of the model
@@ -135,10 +136,10 @@ class KilnBaseModel(BaseModel):
         """
         if isinstance(path, str):
             path = Path(path)
-        cached_model = ModelCache.shared().get_model(path, cls)
+        cached_model = ModelCache.shared().get_model(path, cls, readonly=readonly)
         if cached_model is not None:
             return cached_model
-        with open(path, "r") as file:
+        with open(path, "r", encoding="utf-8") as file:
             # modified time of file for cache invalidation. From file descriptor so it's atomic w read.
             mtime_ns = os.fstat(file.fileno()).st_mtime_ns
             file_data = file.read()
@@ -168,13 +169,20 @@ class KilnBaseModel(BaseModel):
         # Two methods of indicated it's loaded from file:
         # 1) info.context.get("loading_from_file") -> During actual loading, before we can set _loaded_from_file
         # 2) self._loaded_from_file -> After loading, set by the loader
+        if self.loading_from_file(info):
+            return True
+        return self._loaded_from_file
+
+    # indicates the model is currently being loaded from file (not mutating it after)
+    def loading_from_file(self, info: ValidationInfo | None = None) -> bool:
+        # info.context.get("loading_from_file") -> During actual loading, before we can set _loaded_from_file
         if (
             info is not None
             and info.context is not None
             and info.context.get("loading_from_file", False)
         ):
             return True
-        return self._loaded_from_file
+        return False
 
     def save_to_file(self) -> None:
         """Save the model instance to a file.
@@ -190,7 +198,7 @@ class KilnBaseModel(BaseModel):
             )
         path.parent.mkdir(parents=True, exist_ok=True)
         json_data = self.model_dump_json(indent=2, exclude={"path"})
-        with open(path, "w") as file:
+        with open(path, "w", encoding="utf-8") as file:
             file.write(json_data)
         # save the path so even if something like name changes, the file doesn't move
         self.path = path
@@ -342,16 +350,28 @@ class KilnParentedModel(KilnBaseModel, metaclass=ABCMeta):
             return []
 
         # Collect all /relationship/{id}/{base_filename.kiln} files in the relationship folder
-        for child_file in relationship_folder.glob(f"**/{cls.base_filename()}"):
-            yield child_file
+        # manual code instead of glob for performance (5x speedup over glob)
+
+        base_filename = cls.base_filename()
+        # Iterate through immediate subdirectories using scandir for better performance
+        # Benchmark: scandir is 10x faster than glob, so worth the extra code
+        with os.scandir(relationship_folder) as entries:
+            for entry in entries:
+                if not entry.is_dir():
+                    continue
+
+                child_file = Path(entry.path) / base_filename
+                if child_file.is_file():
+                    yield child_file
 
     @classmethod
     def all_children_of_parent_path(
-        cls: Type[PT], parent_path: Path | None
+        cls: Type[PT], parent_path: Path | None, readonly: bool = False
     ) -> list[PT]:
         children = []
         for child_path in cls.iterate_children_paths_of_parent_path(parent_path):
-            children.append(cls.load_from_file(child_path))
+            item = cls.load_from_file(child_path, readonly=readonly)
+            children.append(item)
         return children
 
     @classmethod
@@ -394,8 +414,8 @@ class KilnParentModel(KilnBaseModel, metaclass=ABCMeta):
     def _create_child_method(
         cls, relationship_name: str, child_class: Type[KilnParentedModel]
     ):
-        def child_method(self) -> list[child_class]:
-            return child_class.all_children_of_parent_path(self.path)
+        def child_method(self, readonly: bool = False) -> list[child_class]:
+            return child_class.all_children_of_parent_path(self.path, readonly=readonly)
 
         child_method.__name__ = relationship_name
         child_method.__annotations__ = {"return": List[child_class]}

kiln_ai/datamodel/datamodel_enums.py

@@ -0,0 +1,58 @@
+from enum import Enum, IntEnum
+
+
+class Priority(IntEnum):
+    """Defines priority levels for tasks and requirements, where P0 is highest priority."""
+
+    p0 = 0
+    p1 = 1
+    p2 = 2
+    p3 = 3
+
+
+# Only one rating type for now, but this allows for extensibility if we want to add more in the future
+class TaskOutputRatingType(str, Enum):
+    """Defines the types of rating systems available for task outputs."""
+
+    five_star = "five_star"
+    pass_fail = "pass_fail"
+    pass_fail_critical = "pass_fail_critical"
+    custom = "custom"
+
+
+class StructuredOutputMode(str, Enum):
+    """
+    Enumeration of supported structured output modes.
+
+    - default: let the adapter decide
+    - json_schema: request json using API capabilities for json_schema
+    - function_calling: request json using API capabilities for function calling
+    - json_mode: request json using API's JSON mode, which should return valid JSON, but isn't checking/passing the schema
+    - json_instructions: append instructions to the prompt to request json matching the schema. No API capabilities are used. You should have a custom parser on these models as they will be returning strings.
+    - json_instruction_and_object: append instructions to the prompt to request json matching the schema. Also request the response as json_mode via API capabilities (returning dictionaries).
+    """
+
+    default = "default"
+    json_schema = "json_schema"
+    function_calling_weak = "function_calling_weak"
+    function_calling = "function_calling"
+    json_mode = "json_mode"
+    json_instructions = "json_instructions"
+    json_instruction_and_object = "json_instruction_and_object"
+
+
+class FineTuneStatusType(str, Enum):
+    """
+    The status type of a fine-tune (running, completed, failed, etc).
+    """
+
+    unknown = "unknown"  # server error
+    pending = "pending"
+    running = "running"
+    completed = "completed"
+    failed = "failed"
+
+
+class FinetuneDataStrategy(str, Enum):
+    final_only = "final_only"
+    final_and_intermediate = "final_and_intermediate"

kiln_ai/datamodel/dataset_filters.py

@@ -0,0 +1,114 @@
+from enum import Enum
+from typing import Annotated, Protocol
+
+from pydantic import AfterValidator
+
+from kiln_ai.datamodel.task_run import TaskRun
+
+
+class DatasetFilter(Protocol):
+    """A protocol defining the interface for dataset filters.
+
+    This allows both stateless function-based filters and stateful class-based filters
+    to be used interchangeably, as long as they implement the __call__ method.
+    """
+
+    def __call__(self, task_run: TaskRun) -> bool:
+        """Return True if the task run should be included in the dataset."""
+        ...
+
+
+def AllDatasetFilter(_: TaskRun) -> bool:
+    return True
+
+
+def HighRatingDatasetFilter(task_run: TaskRun) -> bool:
+    if task_run.output is None:
+        return False
+    if task_run.repaired_output is not None:
+        # Repairs always considered high quality
+        return True
+    if task_run.output.rating is None:
+        return False
+    return task_run.output.rating.is_high_quality()
+
+
+def ThinkingModelDatasetFilter(task_run: TaskRun) -> bool:
+    """
+    A filter that returns True if the task has intermediate outputs we can training a 'thinking' model on (reasoning or chain of thought)
+    """
+    return task_run.has_thinking_training_data()
+
+
+def ThinkingModelHighRatedFilter(task_run: TaskRun) -> bool:
+    """
+    A filter that returns True if the task has thinking data and the output is high quality
+    """
+    return ThinkingModelDatasetFilter(task_run) and HighRatingDatasetFilter(task_run)
+
+
+class TagFilter:
+    """
+    A filter that returns True if the task has a tag matching the given tag.
+    """
+
+    def __init__(self, tag: str):
+        self.tag = tag
+
+    def __call__(self, task_run: TaskRun) -> bool:
+        return self.tag in task_run.tags
+
+
+class StaticDatasetFilters(str, Enum):
+    """Dataset filter names."""
+
+    ALL = "all"
+    HIGH_RATING = "high_rating"
+    THINKING_MODEL = "thinking_model"
+    THINKING_MODEL_HIGH_RATED = "thinking_model_high_rated"
+
+
+static_dataset_filters = {
+    StaticDatasetFilters.ALL: AllDatasetFilter,
+    StaticDatasetFilters.HIGH_RATING: HighRatingDatasetFilter,
+    StaticDatasetFilters.THINKING_MODEL: ThinkingModelDatasetFilter,
+    StaticDatasetFilters.THINKING_MODEL_HIGH_RATED: ThinkingModelHighRatedFilter,
+}
+
+DatasetFilterId = Annotated[
+    str,
+    AfterValidator(lambda v: _check_dataset_filter_id(v)),
+]
+"""
+A pydantic type that validates strings containing a valid dataset filter ID.
+
+Dataset filter IDs can be one of:
+- A built-in dataset filter name
+- A tag::<tag> filter, where <tag> is a string
+"""
+
+
+def _check_dataset_filter_id(id: str) -> str:
+    """
+    Check that the dataset filter ID is valid.
+    """
+    if id in static_dataset_filters:
+        return id
+
+    if id.startswith("tag::") and len(id) > 5:
+        return id
+
+    raise ValueError(f"Invalid dataset filter ID: {id}")
+
+
+def dataset_filter_from_id(id: DatasetFilterId) -> DatasetFilter:
+    """
+    Get a dataset filter from an ID.
+    """
+    if id.startswith("tag::") and len(id) > 5:
+        return TagFilter(id[5:])
+
+    if id in static_dataset_filters:
+        return static_dataset_filters[id]
+
+    raise ValueError(f"Invalid dataset filter ID: {id}")

kiln_ai/datamodel/dataset_split.py

@@ -0,0 +1,170 @@
+"""
+Tools for splitting datasets into train/test/validation splits. Includes filters for selecting which task runs to include in each split.
+"""
+
+import math
+import random
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field, model_validator
+
+from kiln_ai.datamodel.basemodel import NAME_FIELD, KilnParentedModel
+from kiln_ai.datamodel.dataset_filters import (
+    DatasetFilter,
+    DatasetFilterId,
+    dataset_filter_from_id,
+)
+
+if TYPE_CHECKING:
+    from kiln_ai.datamodel.task import Task
+
+
+class DatasetSplitDefinition(BaseModel):
+    """
+    A definition of a split in a dataset.
+
+    Example: name="train", description="The training set", percentage=0.8 (80% of the dataset)
+    """
+
+    name: str = NAME_FIELD
+    description: str | None = Field(
+        default=None,
+        description="A description of the dataset for you and your team. Not used in training.",
+    )
+    percentage: float = Field(
+        ge=0.0,
+        le=1.0,
+        description="The percentage of the dataset that this split represents (between 0 and 1).",
+    )
+
+
+AllSplitDefinition: list[DatasetSplitDefinition] = [
+    DatasetSplitDefinition(name="all", percentage=1.0)
+]
+Train80Test20SplitDefinition: list[DatasetSplitDefinition] = [
+    DatasetSplitDefinition(name="train", percentage=0.8),
+    DatasetSplitDefinition(name="test", percentage=0.2),
+]
+Train60Test20Val20SplitDefinition: list[DatasetSplitDefinition] = [
+    DatasetSplitDefinition(name="train", percentage=0.6),
+    DatasetSplitDefinition(name="test", percentage=0.2),
+    DatasetSplitDefinition(name="val", percentage=0.2),
+]
+Train80Test10Val10SplitDefinition: list[DatasetSplitDefinition] = [
+    DatasetSplitDefinition(name="train", percentage=0.8),
+    DatasetSplitDefinition(name="test", percentage=0.1),
+    DatasetSplitDefinition(name="val", percentage=0.1),
+]
+
+
+class DatasetSplit(KilnParentedModel):
+    """
+    A collection of task runs, with optional splits (train, test, validation).
+
+    Used to freeze a dataset into train/test/validation splits for repeatable fine-tuning or other tasks.
+
+    Maintains a list of IDs for each split, to avoid data duplication.
+    """
+
+    name: str = NAME_FIELD
+    description: str | None = Field(
+        default=None,
+        description="A description of the dataset for you and your team. Not used in training.",
+    )
+    splits: list[DatasetSplitDefinition] = Field(
+        default_factory=list,
+        description="The splits in the dataset.",
+    )
+    split_contents: dict[str, list[str]] = Field(
+        description="The contents of each split in the dataset. The key is the split name, and the value is a list of task run IDs.",
+    )
+    filter: DatasetFilterId | None = Field(
+        default=None,
+        description="The filter used to build the dataset.",
+    )
+
+    @model_validator(mode="after")
+    def validate_split_percentages(self) -> "DatasetSplit":
+        total = sum(split.percentage for split in self.splits)
+        if not math.isclose(total, 1.0, rel_tol=1e-9):
+            raise ValueError(f"The sum of split percentages must be 1.0 (got {total})")
+        return self
+
+    @classmethod
+    def from_task(
+        cls,
+        name: str,
+        task: "Task",
+        splits: list[DatasetSplitDefinition],
+        filter_id: DatasetFilterId = "all",
+        description: str | None = None,
+    ):
+        """
+        Build a dataset split from a task.
+        """
+        filter = dataset_filter_from_id(filter_id)
+        split_contents = cls.build_split_contents(task, splits, filter)
+        return cls(
+            parent=task,
+            name=name,
+            description=description,
+            splits=splits,
+            split_contents=split_contents,
+            filter=filter_id,
+        )
+
+    @classmethod
+    def build_split_contents(
+        cls,
+        task: "Task",
+        splits: list[DatasetSplitDefinition],
+        filter: DatasetFilter,
+    ) -> dict[str, list[str]]:
+        valid_ids = []
+        for task_run in task.runs():
+            if filter(task_run):
+                valid_ids.append(task_run.id)
+
+        # Shuffle and split by split percentage
+        random.shuffle(valid_ids)
+        split_contents = {}
+        start_idx = 0
+        remaining_items = len(valid_ids)
+
+        # Handle all splits except the last one
+        for split in splits[:-1]:
+            split_size = round(len(valid_ids) * split.percentage)
+            split_contents[split.name] = valid_ids[start_idx : start_idx + split_size]
+            start_idx += split_size
+            remaining_items -= split_size
+
+        # Last split gets all remaining items (for rounding)
+        if splits:
+            split_contents[splits[-1].name] = valid_ids[start_idx:]
+
+        return split_contents
+
+    def parent_task(self) -> "Task | None":
+        # inline import to avoid circular import
+        from kiln_ai.datamodel import Task
+
+        if not isinstance(self.parent, Task):
+            return None
+        return self.parent
+
+    def missing_count(self) -> int:
+        """
+        Returns:
+            int: the number of task runs that have an ID persisted in this dataset split, but no longer exist in the dataset
+        """
+        parent = self.parent_task()
+        if parent is None:
+            raise ValueError("DatasetSplit has no parent task")
+
+        runs = parent.runs(readonly=True)
+        all_ids = set(run.id for run in runs)
+        all_ids_in_splits = set()
+        for ids in self.split_contents.values():
+            all_ids_in_splits.update(ids)
+        missing = all_ids_in_splits - all_ids
+        return len(missing)