Merge pull request #194 from Kiln-AI/eval_refactor_datamodel

Refactor datamodel
Kiln-AI · Feb 12, 2025 · 616b4c2 · 616b4c2
2 parents 04ba0fa + 75faf43
commit 616b4c2
Show file tree

Hide file tree

Showing 20 changed files with 1,077 additions and 976 deletions.
diff --git a/app/desktop/desktop_server.py b/app/desktop/desktop_server.py
@@ -3,11 +3,10 @@
 import time
 from contextlib import asynccontextmanager
 
+import kiln_ai.datamodel.strict_mode as datamodel_strict_mode
 import kiln_server.server as kiln_server
 import uvicorn
 from fastapi import FastAPI
-from kiln_ai.datamodel import set_strict_mode as set_strict_mode_datamodel
-from kiln_ai.datamodel import strict_mode as strict_mode_datamodel
 
 from app.desktop.studio_server.data_gen_api import connect_data_gen_api
 from app.desktop.studio_server.finetune_api import connect_fine_tune_api
@@ -20,12 +19,12 @@
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Set strict mode on startup
-    original_strict_mode = strict_mode_datamodel()
-    set_strict_mode_datamodel(True)
+    # Set datamodel strict mode on startup
+    original_strict_mode = datamodel_strict_mode.strict_mode()
+    datamodel_strict_mode.set_strict_mode(True)
     yield
-    # Reset strict mode on shutdown
-    set_strict_mode_datamodel(original_strict_mode)
+    # Reset datamodel strict mode on shutdown
+    datamodel_strict_mode.set_strict_mode(original_strict_mode)
 
 
 def make_app():

diff --git a/app/desktop/studio_server/finetune_api.py b/app/desktop/studio_server/finetune_api.py
@@ -18,13 +18,15 @@
     provider_name_from_id,
 )
 from kiln_ai.datamodel import (
-    AllSplitDefinition,
-    DatasetFilterType,
     DatasetSplit,
     Finetune,
     FinetuneDataStrategy,
     FineTuneStatusType,
     Task,
+)
+from kiln_ai.datamodel.dataset_split import (
+    AllSplitDefinition,
+    DatasetFilterType,
     Train60Test20Val20SplitDefinition,
     Train80Test10Val10SplitDefinition,
     Train80Test20SplitDefinition,

diff --git a/app/desktop/studio_server/test_finetune_api.py b/app/desktop/studio_server/test_finetune_api.py
@@ -9,14 +9,16 @@
 from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat
 from kiln_ai.adapters.ml_model_list import KilnModel, KilnModelProvider
 from kiln_ai.datamodel import (
-    AllDatasetFilter,
-    AllSplitDefinition,
     DatasetSplit,
     Finetune,
     FinetuneDataStrategy,
-    HighRatingDatasetFilter,
     Project,
     Task,
+)
+from kiln_ai.datamodel.dataset_split import (
+    AllDatasetFilter,
+    AllSplitDefinition,
+    HighRatingDatasetFilter,
     ThinkingModelDatasetFilter,
     ThinkingModelHighRatedFilter,
     Train60Test20Val20SplitDefinition,
@@ -302,7 +304,7 @@ def test_api_split_types_mapping():
 
 
 def test_api_filter_types_mapping():
-    from kiln_ai.datamodel import dataset_filters
+    from kiln_ai.datamodel.dataset_split import dataset_filters
 
     assert dataset_filters[DatasetFilterType.ALL] == AllDatasetFilter
     assert dataset_filters[DatasetFilterType.HIGH_RATING] == HighRatingDatasetFilter

diff --git a/app/desktop/studio_server/test_server.py b/app/desktop/studio_server/test_server.py
@@ -6,7 +6,7 @@
 import requests
 from fastapi import HTTPException
 from fastapi.testclient import TestClient
-from kiln_ai.datamodel import strict_mode
+from kiln_ai.datamodel.strict_mode import strict_mode
 
 from app.desktop.desktop_server import make_app
 from app.desktop.studio_server.webhost import HTMLStaticFiles

diff --git a/libs/core/kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py b/libs/core/kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py
@@ -16,9 +16,9 @@
     FinetuneDataStrategy,
     StructuredOutputMode,
     Task,
-    Train80Test20SplitDefinition,
 )
 from kiln_ai.datamodel import Finetune as FinetuneModel
+from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
 from kiln_ai.utils.config import Config
 
 

diff --git a/libs/core/kiln_ai/adapters/fine_tune/test_openai_finetune.py b/libs/core/kiln_ai/adapters/fine_tune/test_openai_finetune.py
@@ -15,9 +15,9 @@
     FinetuneDataStrategy,
     StructuredOutputMode,
     Task,
-    Train80Test20SplitDefinition,
 )
 from kiln_ai.datamodel import Finetune as FinetuneModel
+from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
 from kiln_ai.utils.config import Config