kiln-ai 0.13.0__tar.gz → 0.14.0__tar.gz

{kiln_ai-0.13.0 → kiln_ai-0.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiln-ai
-Version: 0.13.0
+Version: 0.14.0
 Summary: Kiln AI
 Project-URL: Homepage, https://getkiln.ai
 Project-URL: Repository, https://github.com/Kiln-AI/kiln
@@ -26,7 +26,7 @@ Requires-Dist: pydantic>=2.9.2
 Requires-Dist: pytest-benchmark>=5.1.0
 Requires-Dist: pytest-cov>=6.0.0
 Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: together>=1.4.6
+Requires-Dist: together
 Requires-Dist: typing-extensions>=4.12.2
 Description-Content-Type: text/markdown
 

{kiln_ai-0.13.0 → kiln_ai-0.14.0}/kiln_ai/adapters/adapter_registry.py

@@ -108,6 +108,10 @@ def adapter_for_task(
                     # 1. To use the correct base URL
                     # 2. We use Ollama's OpenAI compatible API (/v1), and don't just let litellm use the Ollama API. We use more advanced features like json_schema.
                     base_url=ollama_base_url + "/v1",
+                    additional_body_options={
+                        # LiteLLM errors without an api_key, even though Ollama doesn't support one.
+                        "api_key": "NA",
+                    },
                 ),
             )
         case ModelProviderName.fireworks_ai:

{kiln_ai-0.13.0 → kiln_ai-0.14.0}/kiln_ai/adapters/eval/g_eval.py

@@ -297,9 +297,12 @@ The model produced the following output for the task:
 
         total_score = 0.0
         total_probability = 0.0
+        top_logprobs_contains_primary_token = False
 
-        # Process all valid scoring tokens
+        # Process all valid scoring tokens from alternatives
         for top_logprob in token_logprob.top_logprobs:
+            if top_logprob.token == token_logprob.token:
+                top_logprobs_contains_primary_token = True
             token_score = self.score_from_token_string(top_logprob.token)
             if token_score is not None:
                 # Convert logprob to probability
@@ -307,9 +310,21 @@ The model produced the following output for the task:
                 total_score += token_score * probability
                 total_probability += probability
 
+        # Weird OpenAI 4o bug - sometimes the primary token is included in the top logprobs, sometimes not.
+        # Add the primary token back in if excluded
+        if not top_logprobs_contains_primary_token:
+            if token_logprob.logprob == -9999.0:
+                # Another "bug" - sometimes the logprob is -9999.0. This seems to happen when the rest of the logprobs are tiny probability.
+                total_score += primary_token_score * 1.0
+                total_probability += 1.0
+            else:
+                probability = math.exp(token_logprob.logprob)
+                total_score += primary_token_score * probability
+                total_probability += probability
+
         if total_probability <= 0.0:
             raise RuntimeError(
-                f"No valid scoring tokens found for {token_logprob.token}. This should never happen. Please file a bug if you see this."
+                f"No valid scoring tokens found for {token_logprob.token}. This should never happen as the token has a valid score (so it must be excluded from top logprobs). Please file a bug if you see this."
             )
 
         # Normalize by total probability of valid tokens (LLM may have wanted to generate other non-rating tokens, these shouldn't lower score of rating tokens)

{kiln_ai-0.13.0 → kiln_ai-0.14.0}/kiln_ai/adapters/eval/test_g_eval.py

@@ -393,12 +393,13 @@ def test_rating_token_to_score(test_eval_config, test_run_config):
             self.logprob = logprob
 
     class MockTokenLogprob:
-        def __init__(self, token, top_logprobs):
+        def __init__(self, token, top_logprobs, logprob):
             self.token = token
             self.top_logprobs = [MockTopLogprob(t, lp) for t, lp in top_logprobs]
+            self.logprob = logprob
 
     # Test single token case
-    token_logprob = MockTokenLogprob("5", [("5", 0.0)])  # log(1) = 0
+    token_logprob = MockTokenLogprob("5", [("5", 0.0)], logprob=1e-8)  # log(1) = 0
     score = g_eval.rating_token_to_score(token_logprob)
     assert score == 5.0
 
@@ -409,18 +410,22 @@ def test_rating_token_to_score(test_eval_config, test_run_config):
             ("4", math.log(0.6)),  # 60% probability
             ("5", math.log(0.4)),  # 40% probability
         ],
+        logprob=math.log(0.6),
     )
     score = g_eval.rating_token_to_score(token_logprob)
     assert pytest.approx(score) == 4.4  # (4 * 0.6 + 5 * 0.4)
 
     # Test invalid token
-    token_logprob = MockTokenLogprob(":", [(":", 0.0)])
+    token_logprob = MockTokenLogprob(":", [(":", 0.0)], logprob=1e-8)
     assert g_eval.rating_token_to_score(token_logprob) is None
 
-    # Test no valid scoring tokens
-    token_logprob = MockTokenLogprob("5", [])
-    with pytest.raises(RuntimeError, match="No valid scoring tokens found"):
-        g_eval.rating_token_to_score(token_logprob)
+    # Test missing from top logprobs
+    token_logprob = MockTokenLogprob("5", [], logprob=1e-8)
+    assert pytest.approx(g_eval.rating_token_to_score(token_logprob)) == 5.0
+
+    # Test missing from top logprobs, with special case logprob
+    token_logprob = MockTokenLogprob("5", [], logprob=-9999)
+    assert pytest.approx(g_eval.rating_token_to_score(token_logprob)) == 5.0
 
 
 def test_g_eval_system_instruction():

{kiln_ai-0.13.0 → kiln_ai-0.14.0}/kiln_ai/adapters/fine_tune/base_finetune.py

@@ -72,8 +72,6 @@ class BaseFinetuneAdapter(ABC):
         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")
 
@@ -184,21 +182,3 @@ class BaseFinetuneAdapter(ABC):
         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-0.13.0 → kiln_ai-0.14.0}/kiln_ai/adapters/fine_tune/fireworks_finetune.py

@@ -1,4 +1,5 @@
-from typing import Tuple
+import logging
+from typing import List, Tuple
 from uuid import uuid4
 
 import httpx
@@ -13,6 +14,14 @@ from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetF
 from kiln_ai.datamodel import DatasetSplit, StructuredOutputMode, Task
 from kiln_ai.utils.config import Config
 
+logger = logging.getLogger(__name__)
+
+# https://docs.fireworks.ai/fine-tuning/fine-tuning-models#supported-base-models-loras-on-serverless
+serverless_models = [
+    "accounts/fireworks/models/llama-v3p1-8b-instruct",
+    "accounts/fireworks/models/llama-v3p1-70b-instruct",
+]
+
 
 class FireworksFinetune(BaseFinetuneAdapter):
     """
@@ -132,11 +141,18 @@ class FireworksFinetune(BaseFinetuneAdapter):
                 :60
             ]
         )
-        payload = {
+        payload: dict[str, str | dict[str, str | bool]] = {
             "dataset": f"accounts/{account_id}/datasets/{train_file_id}",
             "displayName": display_name,
             "baseModel": self.datamodel.base_model_id,
         }
+        # Add W&B config if API key is set
+        if Config.shared().wandb_api_key:
+            payload["wandbConfig"] = {
+                "enabled": True,
+                "project": "Kiln_AI",
+                "apiKey": Config.shared().wandb_api_key,
+            }
         hyperparameters = self.create_payload_parameters(self.datamodel.parameters)
         payload.update(hyperparameters)
         headers = {
@@ -276,32 +292,54 @@ class FireworksFinetune(BaseFinetuneAdapter):
         return {k: v for k, v in payload.items() if v is not None}
 
     async def _deploy(self) -> bool:
-        # Now we "deploy" the model using PEFT serverless.
-        # A bit complicated: most fireworks deploys are server based.
-        # However, a Lora can be serverless (PEFT).
-        # By calling the deploy endpoint WITHOUT first creating a deployment ID, it will only deploy if it can be done serverless.
-        # https://docs.fireworks.ai/models/deploying#deploying-to-serverless
-        # This endpoint will return 400 if already deployed with code 9, so we consider that a success.
+        if self.datamodel.base_model_id in serverless_models:
+            return await self._deploy_serverless()
+        else:
+            return await self._check_or_deploy_server()
 
+    def api_key_and_account_id(self) -> Tuple[str, str]:
         api_key = Config.shared().fireworks_api_key
         account_id = Config.shared().fireworks_account_id
         if not api_key or not account_id:
             raise ValueError("Fireworks API key or account ID not set")
+        return api_key, account_id
+
+    def deployment_display_name(self) -> str:
+        # Limit the display name to 60 characters
+        display_name = f"Kiln AI fine-tuned model [ID:{self.datamodel.id}][name:{self.datamodel.name}]"[
+            :60
+        ]
+        return display_name
 
+    async def model_id_checking_status(self) -> str | None:
         # 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
+            return None
         if not model_id or not isinstance(model_id, str):
-            return False
+            return None
+        return model_id
+
+    async def _deploy_serverless(self) -> bool:
+        # Now we "deploy" the model using PEFT serverless.
+        # A bit complicated: most fireworks deploys are server based.
+        # However, a Lora can be serverless (PEFT).
+        # By calling the deploy endpoint WITHOUT first creating a deployment ID, it will only deploy if it can be done serverless.
+        # https://docs.fireworks.ai/models/deploying#deploying-to-serverless
+        # This endpoint will return 400 if already deployed with code 9, so we consider that a success.
+
+        api_key, account_id = self.api_key_and_account_id()
 
         url = f"https://api.fireworks.ai/v1/accounts/{account_id}/deployedModels"
-        # Limit the display name to 60 characters
-        display_name = f"Kiln AI fine-tuned model [ID:{self.datamodel.id}][name:{self.datamodel.name}]"[
-            :60
-        ]
+        model_id = await self.model_id_checking_status()
+        if not model_id:
+            logger.error(
+                "Model ID not found - can't deploy model to Fireworks serverless"
+            )
+            return False
+
         payload = {
-            "displayName": display_name,
+            "displayName": self.deployment_display_name(),
             "model": model_id,
         }
         headers = {
@@ -320,4 +358,120 @@ class FireworksFinetune(BaseFinetuneAdapter):
                     self.datamodel.save_to_file()
             return True
 
+        logger.error(
+            f"Failed to deploy model to Fireworks serverless: [{response.status_code}] {response.text}"
+        )
         return False
+
+    async def _check_or_deploy_server(self) -> bool:
+        """
+        Check if the model is already deployed. If not, deploy it to a dedicated server.
+        """
+
+        # Check if the model is already deployed
+        # If it's fine_tune_model_id is set, it might be deployed. However, Fireworks deletes them over time so we need to check.
+        if self.datamodel.fine_tune_model_id:
+            deployments = await self._fetch_all_deployments()
+            for deployment in deployments:
+                if deployment[
+                    "baseModel"
+                ] == self.datamodel.fine_tune_model_id and deployment["state"] in [
+                    "READY",
+                    "CREATING",
+                ]:
+                    return True
+
+        # If the model is not deployed, deploy it
+        return await self._deploy_server()
+
+    async def _deploy_server(self) -> bool:
+        # For models that are not serverless, we just need to deploy the model to a server.
+        # We use a scale-to-zero on-demand deployment. If you stop using it, it
+        # will scale to zero and charges will stop.
+        model_id = await self.model_id_checking_status()
+        if not model_id:
+            logger.error("Model ID not found - can't deploy model to Fireworks server")
+            return False
+
+        api_key, account_id = self.api_key_and_account_id()
+        url = f"https://api.fireworks.ai/v1/accounts/{account_id}/deployments"
+
+        payload = {
+            "displayName": self.deployment_display_name(),
+            "description": "Deployed by Kiln AI",
+            # Allow scale to zero
+            "minReplicaCount": 0,
+            "autoscalingPolicy": {
+                "scaleUpWindow": "30s",
+                "scaleDownWindow": "300s",
+                # Scale to zero after 5 minutes of inactivity - this is the minimum allowed
+                "scaleToZeroWindow": "300s",
+            },
+            "baseModel": model_id,
+        }
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        }
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(url, json=payload, headers=headers)
+
+        if response.status_code == 200:
+            basemodel = response.json().get("baseModel")
+            if basemodel is not None and isinstance(basemodel, str):
+                self.datamodel.fine_tune_model_id = basemodel
+                if self.datamodel.path:
+                    self.datamodel.save_to_file()
+                return True
+
+        logger.error(
+            f"Failed to deploy model to Fireworks server: [{response.status_code}] {response.text}"
+        )
+        return False
+
+    async def _fetch_all_deployments(self) -> List[dict]:
+        """
+        Fetch all deployments for an account.
+        """
+        api_key, account_id = self.api_key_and_account_id()
+
+        url = f"https://api.fireworks.ai/v1/accounts/{account_id}/deployments"
+
+        params = {
+            # Note: filter param does not work for baseModel, which would have been ideal, and ideally would have been documented. Instead we'll fetch all and filter.
+            # Max page size
+            "pageSize": 200,
+        }
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+        }
+
+        deployments = []
+
+        # Paginate through all deployments
+        async with httpx.AsyncClient() as client:
+            while True:
+                response = await client.get(url, params=params, headers=headers)
+                json = response.json()
+                if "deployments" not in json or not isinstance(
+                    json["deployments"], list
+                ):
+                    raise ValueError(
+                        f"Invalid response from Fireworks. Expected list of deployments in 'deployments' key: [{response.status_code}] {response.text}"
+                    )
+                deployments.extend(json["deployments"])
+                next_page_token = json.get("nextPageToken")
+                if (
+                    next_page_token
+                    and isinstance(next_page_token, str)
+                    and len(next_page_token) > 0
+                ):
+                    params = {
+                        "pageSize": 200,
+                        "pageToken": next_page_token,
+                    }
+                else:
+                    break
+
+        return deployments

{kiln_ai-0.13.0 → kiln_ai-0.14.0}/kiln_ai/adapters/fine_tune/test_base_finetune.py

@@ -261,15 +261,6 @@ async def test_create_and_start_no_parent_task_path():
         )
 
 
-def test_check_valid_provider_model():
-    MockFinetune.check_valid_provider_model("openai", "gpt-4o-mini-2024-07-18")
-
-    with pytest.raises(
-        ValueError, match="Provider openai with base model gpt-99 is not available"
-    ):
-        MockFinetune.check_valid_provider_model("openai", "gpt-99")
-
-
 async def test_create_and_start_invalid_train_split(mock_dataset):
     # Test with an invalid train split name
     mock_dataset.split_contents = {"valid_train": [], "valid_test": []}