Merge pull request #294 from leonardmq/leonard/fix-qwq-fine-tune-format

fix: QwQ and R1 finetune format

Author: scosman

app/desktop/studio_server/finetune_api.py

@@ -8,27 +8,26 @@
 from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetFormatter
 from kiln_ai.adapters.fine_tune.finetune_registry import finetune_registry
 from kiln_ai.adapters.ml_model_list import (
+    KilnModel,
+    KilnModelProvider,
+    ModelParserID,
     ModelProviderName,
     built_in_models,
 )
 from kiln_ai.adapters.prompt_builders import (
     chain_of_thought_prompt,
     prompt_builder_from_id,
 )
-from kiln_ai.adapters.provider_tools import (
-    provider_enabled,
-    provider_name_from_id,
-)
+from kiln_ai.adapters.provider_tools import provider_enabled, provider_name_from_id
 from kiln_ai.datamodel import (
     DatasetSplit,
     Finetune,
     FinetuneDataStrategy,
     FineTuneStatusType,
     Task,
 )
-from kiln_ai.datamodel.dataset_filters import (
-    DatasetFilterId,
-)
+from kiln_ai.datamodel.datamodel_enums import THINKING_DATA_STRATEGIES
+from kiln_ai.datamodel.dataset_filters import DatasetFilterId
 from kiln_ai.datamodel.dataset_split import (
     AllSplitDefinition,
     Train60Test20Val20SplitDefinition,
@@ -38,7 +37,7 @@
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.name_generator import generate_memorable_name
 from kiln_server.task_api import task_from_id
-from pydantic import BaseModel
+from pydantic import BaseModel, Field, model_validator
 
 logger = logging.getLogger(__name__)
 
@@ -48,6 +47,12 @@ class FinetuneProviderModel(BaseModel):
 
     name: str
     id: str
+    data_strategies_supported: list[FinetuneDataStrategy] = Field(
+        default_factory=lambda: [
+            FinetuneDataStrategy.final_only,
+            FinetuneDataStrategy.final_and_intermediate,
+        ]
+    )
 
 
 class FinetuneProvider(BaseModel):
@@ -101,6 +106,16 @@ class CreateFinetuneRequest(BaseModel):
     custom_thinking_instructions: str | None = None
     data_strategy: FinetuneDataStrategy
 
+    @model_validator(mode="after")
+    def validate_data_strategy(self) -> "CreateFinetuneRequest":
+        if self.data_strategy not in infer_data_strategies_for_model(
+            built_in_models, self.base_model_id, self.provider
+        ):
+            raise ValueError(
+                f"The data strategy {self.data_strategy} is not supported for the provider model {self.base_model_id}"
+            )
+        return self
+
 
 class FinetuneWithStatus(BaseModel):
     """Finetune with status"""
@@ -198,7 +213,8 @@ async def finetune_providers() -> list[FinetuneProvider]:
                         provider_models[provider.name] = []
                     provider_models[provider.name].append(
                         FinetuneProviderModel(
-                            name=model.friendly_name, id=provider.provider_finetune_id
+                            name=model.friendly_name,
+                            id=provider.provider_finetune_id,
                         )
                     )
 
@@ -212,14 +228,19 @@ async def finetune_providers() -> list[FinetuneProvider]:
         # Create provider entries
         providers: list[FinetuneProvider] = []
         for provider_name, models in provider_models.items():
-            providers.append(
-                FinetuneProvider(
-                    name=provider_name_from_id(provider_name),
-                    id=provider_name,
-                    enabled=await provider_enabled(provider_name),
-                    models=models,
+            # attach the compatible data strategies to each model
+            for model in models:
+                model.data_strategies_supported = infer_data_strategies_for_model(
+                    built_in_models, model.id, provider_name
                 )
+
+            provider = FinetuneProvider(
+                name=provider_name_from_id(provider_name),
+                id=provider_name,
+                enabled=await provider_enabled(provider_name),
+                models=models,
             )
+            providers.append(provider)
 
         return providers
 
@@ -326,6 +347,7 @@ async def download_dataset_jsonl(
                 status_code=400,
                 detail=f"Data strategy '{data_strategy}' not found",
             )
+
         data_strategy_typed = FinetuneDataStrategy(data_strategy)
 
         task = task_from_id(project_id, task_id)
@@ -406,10 +428,13 @@ def thinking_instructions_from_request(
     data_strategy: FinetuneDataStrategy,
     custom_thinking_instructions: str | None,
 ) -> str | None:
-    if data_strategy != FinetuneDataStrategy.final_and_intermediate:
+    if data_strategy not in THINKING_DATA_STRATEGIES:
         # Not using COT/Thinking style
         return None
 
+    if data_strategy == FinetuneDataStrategy.final_and_intermediate_r1_compatible:
+        return None
+
     if custom_thinking_instructions:
         # prefer custom instructions
         return custom_thinking_instructions
@@ -477,3 +502,58 @@ async def fetch_fireworks_finetune_models() -> list[FinetuneProviderModel]:
             )
 
     return tuneable_models
+
+
+DEFAULT_DATA_STRATEGIES = [
+    FinetuneDataStrategy.final_only,
+    FinetuneDataStrategy.final_and_intermediate,
+]
+
+
+def data_strategies_from_model_provider(
+    provider: KilnModelProvider,
+) -> list[FinetuneDataStrategy]:
+    if provider.parser == ModelParserID.r1_thinking:
+        return [
+            FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+        ]
+    return DEFAULT_DATA_STRATEGIES
+
+
+def data_strategies_from_finetune_id(
+    provider_finetune_id: str,
+) -> list[FinetuneDataStrategy]:
+    if "qwen3" in provider_finetune_id.lower():
+        return [
+            FinetuneDataStrategy.final_only,
+            FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+        ]
+
+    r1_must_include = ["r1", "qwq"]
+    if any(substring in provider_finetune_id.lower() for substring in r1_must_include):
+        return [
+            FinetuneDataStrategy.final_and_intermediate_r1_compatible,
+        ]
+    return DEFAULT_DATA_STRATEGIES
+
+
+def infer_data_strategies_for_model(
+    available_models: list[KilnModel],
+    provider_finetune_id: str,
+    provider_name: str,
+) -> list[FinetuneDataStrategy]:
+    # we don't have built-in models for fireworks models, so we infer the data strategy from the model name
+    if provider_name == ModelProviderName.fireworks_ai:
+        return data_strategies_from_finetune_id(provider_finetune_id)
+
+    # where we have built-in models, we can infer the data strategy from the object itself
+    for model in available_models:
+        for provider in model.providers:
+            if (
+                provider.name == provider_name
+                and provider.provider_finetune_id == provider_finetune_id
+            ):
+                return data_strategies_from_model_provider(provider)
+
+    # for everything else, we don't know what the data strategy is, so we use the default
+    return DEFAULT_DATA_STRATEGIES